update

README.md

@@ -1,5 +1,5 @@
 ---
-title: NX Denoise
+title: CC Denoise
 emoji: 🐢
 colorFrom: purple
 colorTo: blue
@@ -9,7 +9,7 @@ license: apache-2.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
-## NX Denoise
+## CC Denoise
 
 
 ### datasets

examples/dtln_mp3_to_wav/run.sh

@@ -1,168 +0,0 @@
-#!/usr/bin/env bash
-
-: <<'END'
-
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir-256 --final_model_name dtln-256-nx-dns3 \
---config_file "yaml/config-256.yaml" \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
-
-
-sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir-512 --final_model_name dtln-512-nx-dns3 \
---config_file "yaml/config-512.yaml" \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
-
-
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name dtnl-1024-nx2 --final_model_name dtln-1024-nx2 \
---config_file "yaml/config-1024.yaml" \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/nx-noise" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2"
-
-
-bash run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name dtln-256-nx2-dns3-mp3 --final_model_name dtln-256-nx2-dns3-mp3 \
---config_file "yaml/config-256.yaml" \
---audio_dir "/data/tianxing/HuggingDatasets/nx_noise/data" \
-
-
-END
-
-
-# params
-system_version="windows";
-verbose=true;
-stage=0 # start from 0 if you need to start from data preparation
-stop_stage=9
-
-work_dir="$(pwd)"
-file_folder_name=file_folder_name
-final_model_name=final_model_name
-config_file="yaml/config.yaml"
-limit=10
-
-audio_dir=/data/tianxing/HuggingDatasets/nx_noise/data
-
-max_count=-1
-
-nohup_name=nohup.out
-
-# model params
-batch_size=64
-max_epochs=200
-save_top_k=10
-patience=5
-
-
-# parse options
-while true; do
-  [ -z "${1:-}" ] && break;  # break if there are no arguments
-  case "$1" in
-    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
-      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
-      old_value="(eval echo \\$$name)";
-      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
-        was_bool=true;
-      else
-        was_bool=false;
-      fi
-
-      # Set the variable to the right value-- the escaped quotes make it work if
-      # the option had spaces, like --cmd "queue.pl -sync y"
-      eval "${name}=\"$2\"";
-
-      # Check that Boolean-valued arguments are really Boolean.
-      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
-        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
-        exit 1;
-      fi
-      shift 2;
-      ;;
-
-    *) break;
-  esac
-done
-
-file_dir="${work_dir}/${file_folder_name}"
-final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
-evaluation_audio_dir="${file_dir}/evaluation_audio"
-
-train_dataset="${file_dir}/train.jsonl"
-valid_dataset="${file_dir}/valid.jsonl"
-
-$verbose && echo "system_version: ${system_version}"
-$verbose && echo "file_folder_name: ${file_folder_name}"
-
-if [ $system_version == "windows" ]; then
-  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
-elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
-  #source /data/local/bin/nx_denoise/bin/activate
-  alias python3='/data/local/bin/nx_denoise/bin/python3'
-fi
-
-
-if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
-  $verbose && echo "stage 1: prepare data"
-  cd "${work_dir}" || exit 1
-  python3 step_1_prepare_data.py \
-  --file_dir "${file_dir}" \
-  --audio_dir "${audio_dir}" \
-  --train_dataset "${train_dataset}" \
-  --valid_dataset "${valid_dataset}" \
-  --max_count "${max_count}" \
-
-fi
-
-
-if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
-  $verbose && echo "stage 2: train model"
-  cd "${work_dir}" || exit 1
-  python3 step_2_train_model.py \
-  --train_dataset "${train_dataset}" \
-  --valid_dataset "${valid_dataset}" \
-  --serialization_dir "${file_dir}" \
-  --config_file "${config_file}" \
-
-fi
-
-
-if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
-  $verbose && echo "stage 3: test model"
-  cd "${work_dir}" || exit 1
-  python3 step_3_evaluation.py \
-  --valid_dataset "${valid_dataset}" \
-  --model_dir "${file_dir}/best" \
-  --evaluation_audio_dir "${evaluation_audio_dir}" \
-  --limit "${limit}" \
-
-fi
-
-
-if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
-  $verbose && echo "stage 4: collect files"
-  cd "${work_dir}" || exit 1
-
-  mkdir -p ${final_model_dir}
-
-  cp "${file_dir}/best"/* "${final_model_dir}"
-  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
-
-  cd "${final_model_dir}/.." || exit 1;
-
-  if [ -e "${final_model_name}.zip" ]; then
-    rm -rf "${final_model_name}_backup.zip"
-    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
-  fi
-
-  zip -r "${final_model_name}.zip" "${final_model_name}"
-  rm -rf "${final_model_name}"
-
-fi
-
-
-if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
-  $verbose && echo "stage 5: clear file_dir"
-  cd "${work_dir}" || exit 1
-
-  rm -rf "${file_dir}";
-
-fi

examples/dtln_mp3_to_wav/step_1_prepare_data.py

@@ -1,127 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import argparse
-import json
-import os
-from pathlib import Path
-import random
-import sys
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import librosa
-import numpy as np
-from tqdm import tqdm
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--file_dir", default="./", type=str)
-
-    parser.add_argument(
-        "--audio_dir",
-        default="E:/Users/tianx/HuggingDatasets/nx_noise/data/speech",
-        type=str
-    )
-
-    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
-    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
-
-    parser.add_argument("--duration", default=4.0, type=float)
-
-    parser.add_argument("--target_sample_rate", default=8000, type=int)
-
-    parser.add_argument("--max_count", default=-1, type=int)
-
-    args = parser.parse_args()
-    return args
-
-
-def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000, max_epoch: int = 1):
-    data_dir = Path(data_dir)
-    for epoch_idx in range(max_epoch):
-        for filename in data_dir.glob("**/*.wav"):
-            signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
-            raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
-
-            if raw_duration < duration:
-                # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
-                continue
-            if signal.ndim != 1:
-                raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
-
-            signal_length = len(signal)
-            win_size = int(duration * sample_rate)
-            for begin in range(0, signal_length - win_size, win_size):
-                if np.sum(signal[begin: begin+win_size]) == 0:
-                    continue
-                row = {
-                    "epoch_idx": epoch_idx,
-                    "filename": filename.as_posix(),
-                    "raw_duration": round(raw_duration, 4),
-                    "offset": round(begin / sample_rate, 4),
-                    "duration": round(duration, 4),
-                }
-                yield row
-
-
-def main():
-    args = get_args()
-
-    file_dir = Path(args.file_dir)
-    file_dir.mkdir(exist_ok=True)
-
-    audio_dir = Path(args.audio_dir)
-
-    audio_generator = target_second_signal_generator(
-        audio_dir.as_posix(),
-        duration=args.duration,
-        sample_rate=args.target_sample_rate,
-        max_epoch=1,
-    )
-    count = 0
-    process_bar = tqdm(desc="build dataset jsonl")
-    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
-        for audio in audio_generator:
-            if count >= args.max_count > 0:
-                break
-
-            filename = audio["filename"]
-            raw_duration = audio["raw_duration"]
-            offset = audio["offset"]
-            duration = audio["duration"]
-
-            random1 = random.random()
-            random2 = random.random()
-
-            row = {
-                "count": count,
-
-                "filename": filename,
-                "raw_duration": raw_duration,
-                "offset": offset,
-                "duration": duration,
-
-                "random1": random1,
-            }
-            row = json.dumps(row, ensure_ascii=False)
-            if random2 < (1 / 300):
-                fvalid.write(f"{row}\n")
-            else:
-                ftrain.write(f"{row}\n")
-
-            count += 1
-            duration_seconds = count * args.duration
-            duration_hours = duration_seconds / 3600
-
-            process_bar.update(n=1)
-            process_bar.set_postfix({
-                "duration_hours": round(duration_hours, 4),
-            })
-
-    return
-
-
-if __name__ == "__main__":
-    main()

examples/dtln_mp3_to_wav/step_2_train_model.py

@@ -1,445 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-"""
-https://github.com/breizhn/DTLN
-
-"""
-import argparse
-import json
-import logging
-from logging.handlers import TimedRotatingFileHandler
-import os
-import platform
-from pathlib import Path
-import random
-import sys
-import shutil
-from typing import List
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-from torch.utils.data.dataloader import DataLoader
-from tqdm import tqdm
-
-from toolbox.torch.utils.data.dataset.mp3_to_wav_jsonl_dataset import Mp3ToWavJsonlDataset
-from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
-from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
-from toolbox.torchaudio.metrics.pesq import run_pesq_score
-from toolbox.torchaudio.models.dtln.configuration_dtln import DTLNConfig
-from toolbox.torchaudio.models.dtln.modeling_dtln import DTLNModel, DTLNPretrainedModel
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
-    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
-
-    parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
-    parser.add_argument("--patience", default=30, type=int)
-    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
-
-    parser.add_argument("--config_file", default="config.yaml", type=str)
-
-    args = parser.parse_args()
-    return args
-
-
-def logging_config(file_dir: str):
-    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
-
-    logging.basicConfig(format=fmt,
-                        datefmt="%m/%d/%Y %H:%M:%S",
-                        level=logging.INFO)
-    file_handler = TimedRotatingFileHandler(
-        filename=os.path.join(file_dir, "main.log"),
-        encoding="utf-8",
-        when="D",
-        interval=1,
-        backupCount=7
-    )
-    file_handler.setLevel(logging.INFO)
-    file_handler.setFormatter(logging.Formatter(fmt))
-    logger = logging.getLogger(__name__)
-    logger.addHandler(file_handler)
-
-    return logger
-
-
-class CollateFunction(object):
-    def __init__(self):
-        pass
-
-    def __call__(self, batch: List[dict]):
-        mp3_waveform_list = list()
-        wav_waveform_list = list()
-
-        for sample in batch:
-            mp3_waveform: torch.Tensor = sample["mp3_waveform"]
-            wav_waveform: torch.Tensor = sample["wav_waveform"]
-
-            mp3_waveform_list.append(mp3_waveform)
-            wav_waveform_list.append(wav_waveform)
-
-        mp3_waveform_list = torch.stack(mp3_waveform_list)
-        wav_waveform_list = torch.stack(wav_waveform_list)
-
-        # assert
-        if torch.any(torch.isnan(mp3_waveform_list)) or torch.any(torch.isinf(mp3_waveform_list)):
-            raise AssertionError("nan or inf in mp3_waveform_list")
-        if torch.any(torch.isnan(wav_waveform_list)) or torch.any(torch.isinf(wav_waveform_list)):
-            raise AssertionError("nan or inf in wav_waveform_list")
-
-        return mp3_waveform_list, wav_waveform_list
-
-
-collate_fn = CollateFunction()
-
-
-def main():
-    args = get_args()
-
-    config = DTLNConfig.from_pretrained(
-        pretrained_model_name_or_path=args.config_file,
-    )
-
-    serialization_dir = Path(args.serialization_dir)
-    serialization_dir.mkdir(parents=True, exist_ok=True)
-
-    logger = logging_config(serialization_dir)
-
-    random.seed(config.seed)
-    np.random.seed(config.seed)
-    torch.manual_seed(config.seed)
-    logger.info(f"set seed: {config.seed}")
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    n_gpu = torch.cuda.device_count()
-    logger.info(f"GPU available count: {n_gpu}; device: {device}")
-
-    # datasets
-    train_dataset = Mp3ToWavJsonlDataset(
-        jsonl_file=args.train_dataset,
-        expected_sample_rate=config.sample_rate,
-        max_wave_value=32768.0,
-        # skip=225000,
-    )
-    valid_dataset = Mp3ToWavJsonlDataset(
-        jsonl_file=args.valid_dataset,
-        expected_sample_rate=config.sample_rate,
-        max_wave_value=32768.0,
-    )
-    train_data_loader = DataLoader(
-        dataset=train_dataset,
-        batch_size=config.batch_size,
-        # shuffle=True,
-        sampler=None,
-        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
-        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
-        collate_fn=collate_fn,
-        pin_memory=False,
-        prefetch_factor=None if platform.system() == "Windows" else 2,
-    )
-    valid_data_loader = DataLoader(
-        dataset=valid_dataset,
-        batch_size=config.batch_size,
-        # shuffle=True,
-        sampler=None,
-        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
-        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
-        collate_fn=collate_fn,
-        pin_memory=False,
-        prefetch_factor=None if platform.system() == "Windows" else 2,
-    )
-
-    # models
-    logger.info(f"prepare models. config_file: {args.config_file}")
-    model = DTLNPretrainedModel(config).to(device)
-    model.to(device)
-    model.train()
-
-    # optimizer
-    logger.info("prepare optimizer, lr_scheduler, loss_fn, evaluation_metric")
-    optimizer = torch.optim.AdamW(model.parameters(), config.lr)
-
-    # resume training
-    last_step_idx = -1
-    last_epoch = -1
-    for step_idx_str in serialization_dir.glob("steps-*"):
-        step_idx_str = Path(step_idx_str)
-        step_idx = step_idx_str.stem.split("-")[1]
-        step_idx = int(step_idx)
-        if step_idx > last_step_idx:
-            last_step_idx = step_idx
-    # last_epoch = 1
-
-    if last_step_idx != -1:
-        logger.info(f"resume from steps-{last_step_idx}.")
-        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
-
-        logger.info(f"load state dict for model.")
-        with open(model_pt.as_posix(), "rb") as f:
-            state_dict = torch.load(f, map_location="cpu", weights_only=True)
-        model.load_state_dict(state_dict, strict=True)
-
-    if config.lr_scheduler == "CosineAnnealingLR":
-        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
-            optimizer,
-            last_epoch=last_epoch,
-            # T_max=10 * config.eval_steps,
-            # eta_min=0.01 * config.lr,
-            **config.lr_scheduler_kwargs,
-        )
-    elif config.lr_scheduler == "MultiStepLR":
-        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-            optimizer,
-            last_epoch=last_epoch,
-            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
-        )
-    else:
-        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
-
-    neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
-    mr_stft_loss_fn = MultiResolutionSTFTLoss(
-        fft_size_list=[256, 512, 1024],
-        win_size_list=[256, 512, 1024],
-        hop_size_list=[128, 256, 512],
-        factor_sc=1.5,
-        factor_mag=1.0,
-        reduction="mean"
-    ).to(device)
-    audio_l1_loss_fn = nn.L1Loss(reduction="mean")
-
-    # training loop
-
-    # state
-    average_pesq_score = 1000000000
-    average_loss = 1000000000
-    average_mr_stft_loss = 1000000000
-    average_audio_l1_loss = 1000000000
-    average_neg_si_snr_loss = 1000000000
-
-    model_list = list()
-    best_epoch_idx = None
-    best_step_idx = None
-    best_metric = None
-    patience_count = 0
-
-    step_idx = 0 if last_step_idx == -1 else last_step_idx
-
-    logger.info("training")
-    early_stop_flag = False
-    for epoch_idx in range(max(0, last_epoch+1), config.max_epochs):
-        if early_stop_flag:
-            break
-
-        # train
-        model.train()
-
-        total_pesq_score = 0.
-        total_loss = 0.
-        total_mr_stft_loss = 0.
-        total_audio_l1_loss = 0.
-        total_neg_si_snr_loss = 0.
-        total_batches = 0.
-
-        progress_bar_train = tqdm(
-            initial=step_idx,
-            desc="Training; epoch-{}".format(epoch_idx),
-        )
-        for train_batch in train_data_loader:
-            mp3_audios, wav_audios = train_batch
-            noisy_audios: torch.Tensor = mp3_audios.to(device)
-            clean_audios: torch.Tensor = wav_audios.to(device)
-
-            denoise_audios = model.forward(noisy_audios)
-            denoise_audios = torch.squeeze(denoise_audios, dim=1)
-
-            mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
-            audio_l1_loss = audio_l1_loss_fn.forward(denoise_audios, clean_audios)
-            neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
-
-            loss = 1.0 * mr_stft_loss + 1.0 * audio_l1_loss + 1.0 * neg_si_snr_loss
-            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
-                logger.info(f"find nan or inf in loss.")
-                continue
-
-            denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
-            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
-
-            optimizer.zero_grad()
-            loss.backward()
-            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.clip_grad_norm)
-            optimizer.step()
-            lr_scheduler.step()
-
-            total_pesq_score += pesq_score
-            total_loss += loss.item()
-            total_mr_stft_loss += mr_stft_loss.item()
-            total_audio_l1_loss += audio_l1_loss.item()
-            total_neg_si_snr_loss += neg_si_snr_loss.item()
-            total_batches += 1
-
-            average_pesq_score = round(total_pesq_score / total_batches, 4)
-            average_loss = round(total_loss / total_batches, 4)
-            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
-            average_audio_l1_loss = round(total_audio_l1_loss / total_batches, 4)
-            average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
-
-            progress_bar_train.update(1)
-            progress_bar_train.set_postfix({
-                "lr": lr_scheduler.get_last_lr()[0],
-                "pesq_score": average_pesq_score,
-                "loss": average_loss,
-                "mr_stft_loss": average_mr_stft_loss,
-                "audio_l1_loss": average_audio_l1_loss,
-                "neg_si_snr_loss": average_neg_si_snr_loss,
-            })
-
-            # evaluation
-            step_idx += 1
-            if step_idx % config.eval_steps == 0:
-                model.eval()
-                with torch.no_grad():
-                    torch.cuda.empty_cache()
-
-                    total_pesq_score = 0.
-                    total_loss = 0.
-                    total_mr_stft_loss = 0.
-                    total_audio_l1_loss = 0.
-                    total_neg_si_snr_loss = 0.
-                    total_batches = 0.
-
-                    progress_bar_train.close()
-                    progress_bar_eval = tqdm(
-                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
-                    )
-                    for eval_batch in valid_data_loader:
-                        mp3_audios, wav_audios = eval_batch
-                        noisy_audios: torch.Tensor = mp3_audios.to(device)
-                        clean_audios: torch.Tensor = wav_audios.to(device)
-
-                        denoise_audios = model.forward(noisy_audios)
-                        denoise_audios = torch.squeeze(denoise_audios, dim=1)
-
-                        mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
-                        audio_l1_loss = audio_l1_loss_fn.forward(denoise_audios, clean_audios)
-                        neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
-
-                        loss = 1.0 * mr_stft_loss + 1.0 * audio_l1_loss + 1.0 * neg_si_snr_loss
-                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
-                            logger.info(f"find nan or inf in loss.")
-                            continue
-
-                        denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
-                        clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
-
-                        total_pesq_score += pesq_score
-                        total_loss += loss.item()
-                        total_mr_stft_loss += mr_stft_loss.item()
-                        total_audio_l1_loss += audio_l1_loss.item()
-                        total_neg_si_snr_loss += neg_si_snr_loss.item()
-                        total_batches += 1
-
-                        average_pesq_score = round(total_pesq_score / total_batches, 4)
-                        average_loss = round(total_loss / total_batches, 4)
-                        average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
-                        average_audio_l1_loss = round(total_audio_l1_loss / total_batches, 4)
-                        average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
-
-                        progress_bar_eval.update(1)
-                        progress_bar_eval.set_postfix({
-                            "lr": lr_scheduler.get_last_lr()[0],
-                            "pesq_score": average_pesq_score,
-                            "loss": average_loss,
-                            "mr_stft_loss": average_mr_stft_loss,
-                            "audio_l1_loss": average_audio_l1_loss,
-                            "neg_si_snr_loss": average_neg_si_snr_loss,
-
-                        })
-
-                    total_pesq_score = 0.
-                    total_loss = 0.
-                    total_mr_stft_loss = 0.
-                    total_audio_l1_loss = 0.
-                    total_neg_si_snr_loss = 0.
-                    total_batches = 0.
-
-                    progress_bar_eval.close()
-                    progress_bar_train = tqdm(
-                        initial=progress_bar_train.n,
-                        postfix=progress_bar_train.postfix,
-                        desc=progress_bar_train.desc,
-                    )
-
-                    # save path
-                    save_dir = serialization_dir / "steps-{}".format(step_idx)
-                    save_dir.mkdir(parents=True, exist_ok=False)
-
-                    # save models
-                    model.save_pretrained(save_dir.as_posix())
-
-                    model_list.append(save_dir)
-                    if len(model_list) >= args.num_serialized_models_to_keep:
-                        model_to_delete: Path = model_list.pop(0)
-                        shutil.rmtree(model_to_delete.as_posix())
-
-                    # save metric
-                    if best_metric is None:
-                        best_epoch_idx = epoch_idx
-                        best_step_idx = step_idx
-                        best_metric = average_pesq_score
-                    elif average_pesq_score >= best_metric:
-                        # great is better.
-                        best_epoch_idx = epoch_idx
-                        best_step_idx = step_idx
-                        best_metric = average_pesq_score
-                    else:
-                        pass
-
-                    metrics = {
-                        "epoch_idx": epoch_idx,
-                        "best_epoch_idx": best_epoch_idx,
-                        "best_step_idx": best_step_idx,
-                        "pesq_score": average_pesq_score,
-                        "loss": average_loss,
-                        "mr_stft_loss": average_mr_stft_loss,
-                        "audio_l1_loss": average_audio_l1_loss,
-                        "neg_si_snr_loss": average_neg_si_snr_loss,
-                    }
-                    metrics_filename = save_dir / "metrics_epoch.json"
-                    with open(metrics_filename, "w", encoding="utf-8") as f:
-                        json.dump(metrics, f, indent=4, ensure_ascii=False)
-
-                    # save best
-                    best_dir = serialization_dir / "best"
-                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
-                        if best_dir.exists():
-                            shutil.rmtree(best_dir)
-                        shutil.copytree(save_dir, best_dir)
-
-                    # early stop
-                    early_stop_flag = False
-                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
-                        patience_count = 0
-                    else:
-                        patience_count += 1
-                    if patience_count >= args.patience:
-                        early_stop_flag = True
-
-                    # early stop
-                    if early_stop_flag:
-                        break
-                model.train()
-
-    return
-
-
-if __name__ == "__main__":
-    main()

examples/dtln_mp3_to_wav/yaml/config-1024.yaml

@@ -1,29 +0,0 @@
-model_name: "DTLN"
-
-# spec
-sample_rate: 8000
-fft_size: 512
-hop_size: 128
-win_type: hann
-
-# data
-min_snr_db: -5
-max_snr_db: 25
-
-# model
-encoder_size: 1024
-
-# train
-lr: 0.001
-lr_scheduler: "CosineAnnealingLR"
-lr_scheduler_kwargs:
-  T_max: 250000
-  eta_min: 0.0001
-
-max_epochs: 100
-clip_grad_norm: 10.0
-seed: 1234
-
-num_workers: 4
-batch_size: 64
-eval_steps: 15000

examples/dtln_mp3_to_wav/yaml/config-256.yaml

@@ -1,29 +0,0 @@
-model_name: "DTLN"
-
-# spec
-sample_rate: 8000
-fft_size: 256
-hop_size: 128
-win_type: hann
-
-# data
-min_snr_db: -5
-max_snr_db: 25
-
-# model
-encoder_size: 256
-
-# train
-lr: 0.001
-lr_scheduler: "CosineAnnealingLR"
-lr_scheduler_kwargs:
-  T_max: 250000
-  eta_min: 0.0001
-
-max_epochs: 100
-clip_grad_norm: 10.0
-seed: 1234
-
-num_workers: 4
-batch_size: 64
-eval_steps: 15000

examples/dtln_mp3_to_wav/yaml/config-512.yaml

@@ -1,29 +0,0 @@
-model_name: "DTLN"
-
-# spec
-sample_rate: 8000
-fft_size: 512
-hop_size: 128
-win_type: hann
-
-# data
-min_snr_db: -5
-max_snr_db: 25
-
-# model
-encoder_size: 512
-
-# train
-lr: 0.001
-lr_scheduler: "CosineAnnealingLR"
-lr_scheduler_kwargs:
-  T_max: 250000
-  eta_min: 0.0001
-
-max_epochs: 100
-clip_grad_norm: 10.0
-seed: 1234
-
-num_workers: 4
-batch_size: 64
-eval_steps: 15000

examples/frcrn_mp3_to_wav/run.sh

@@ -1,156 +0,0 @@
-#!/usr/bin/env bash
-
-: <<'END'
-
-
-sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name frcrn-20-512-nx-dns3 \
---config_file "yaml/config-10.yaml" \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech"
-
-
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name frcrn-10-nx2 \
---config_file "yaml/config-10.yaml" \
---audio_dir "/data/tianxing/HuggingDatasets/nx_noise/data" \
-
-END
-
-
-# params
-system_version="windows";
-verbose=true;
-stage=0 # start from 0 if you need to start from data preparation
-stop_stage=9
-
-work_dir="$(pwd)"
-file_folder_name=file_folder_name
-final_model_name=final_model_name
-config_file="yaml/config.yaml"
-limit=10
-
-audio_dir=/data/tianxing/HuggingDatasets/nx_noise/data
-
-max_count=10000000
-
-nohup_name=nohup.out
-
-# model params
-batch_size=64
-max_epochs=200
-save_top_k=10
-patience=5
-
-
-# parse options
-while true; do
-  [ -z "${1:-}" ] && break;  # break if there are no arguments
-  case "$1" in
-    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
-      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
-      old_value="(eval echo \\$$name)";
-      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
-        was_bool=true;
-      else
-        was_bool=false;
-      fi
-
-      # Set the variable to the right value-- the escaped quotes make it work if
-      # the option had spaces, like --cmd "queue.pl -sync y"
-      eval "${name}=\"$2\"";
-
-      # Check that Boolean-valued arguments are really Boolean.
-      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
-        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
-        exit 1;
-      fi
-      shift 2;
-      ;;
-
-    *) break;
-  esac
-done
-
-file_dir="${work_dir}/${file_folder_name}"
-final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
-evaluation_audio_dir="${file_dir}/evaluation_audio"
-
-train_dataset="${file_dir}/train.jsonl"
-valid_dataset="${file_dir}/valid.jsonl"
-
-$verbose && echo "system_version: ${system_version}"
-$verbose && echo "file_folder_name: ${file_folder_name}"
-
-if [ $system_version == "windows" ]; then
-  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
-elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
-  #source /data/local/bin/nx_denoise/bin/activate
-  alias python3='/data/local/bin/nx_denoise/bin/python3'
-fi
-
-
-if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
-  $verbose && echo "stage 1: prepare data"
-  cd "${work_dir}" || exit 1
-  python3 step_1_prepare_data.py \
-  --file_dir "${file_dir}" \
-  --audio_dir "${audio_dir}" \
-  --train_dataset "${train_dataset}" \
-  --valid_dataset "${valid_dataset}" \
-  --max_count "${max_count}" \
-
-fi
-
-
-if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
-  $verbose && echo "stage 2: train model"
-  cd "${work_dir}" || exit 1
-  python3 step_2_train_model.py \
-  --train_dataset "${train_dataset}" \
-  --valid_dataset "${valid_dataset}" \
-  --serialization_dir "${file_dir}" \
-  --config_file "${config_file}" \
-
-fi
-
-
-if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
-  $verbose && echo "stage 3: test model"
-  cd "${work_dir}" || exit 1
-  python3 step_3_evaluation.py \
-  --valid_dataset "${valid_dataset}" \
-  --model_dir "${file_dir}/best" \
-  --evaluation_audio_dir "${evaluation_audio_dir}" \
-  --limit "${limit}" \
-
-fi
-
-
-if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
-  $verbose && echo "stage 4: collect files"
-  cd "${work_dir}" || exit 1
-
-  mkdir -p ${final_model_dir}
-
-  cp "${file_dir}/best"/* "${final_model_dir}"
-  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
-
-  cd "${final_model_dir}/.." || exit 1;
-
-  if [ -e "${final_model_name}.zip" ]; then
-    rm -rf "${final_model_name}_backup.zip"
-    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
-  fi
-
-  zip -r "${final_model_name}.zip" "${final_model_name}"
-  rm -rf "${final_model_name}"
-
-fi
-
-
-if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
-  $verbose && echo "stage 5: clear file_dir"
-  cd "${work_dir}" || exit 1
-
-  rm -rf "${file_dir}";
-
-fi

examples/frcrn_mp3_to_wav/step_1_prepare_data.py

@@ -1,127 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import argparse
-import json
-import os
-from pathlib import Path
-import random
-import sys
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import librosa
-import numpy as np
-from tqdm import tqdm
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--file_dir", default="./", type=str)
-
-    parser.add_argument(
-        "--audio_dir",
-        default="E:/Users/tianx/HuggingDatasets/nx_noise/data/speech",
-        type=str
-    )
-
-    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
-    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
-
-    parser.add_argument("--duration", default=4.0, type=float)
-
-    parser.add_argument("--target_sample_rate", default=8000, type=int)
-
-    parser.add_argument("--max_count", default=-1, type=int)
-
-    args = parser.parse_args()
-    return args
-
-
-def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000, max_epoch: int = 1):
-    data_dir = Path(data_dir)
-    for epoch_idx in range(max_epoch):
-        for filename in data_dir.glob("**/*.wav"):
-            signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
-            raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
-
-            if raw_duration < duration:
-                # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
-                continue
-            if signal.ndim != 1:
-                raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
-
-            signal_length = len(signal)
-            win_size = int(duration * sample_rate)
-            for begin in range(0, signal_length - win_size, win_size):
-                if np.sum(signal[begin: begin+win_size]) == 0:
-                    continue
-                row = {
-                    "epoch_idx": epoch_idx,
-                    "filename": filename.as_posix(),
-                    "raw_duration": round(raw_duration, 4),
-                    "offset": round(begin / sample_rate, 4),
-                    "duration": round(duration, 4),
-                }
-                yield row
-
-
-def main():
-    args = get_args()
-
-    file_dir = Path(args.file_dir)
-    file_dir.mkdir(exist_ok=True)
-
-    audio_dir = Path(args.audio_dir)
-
-    audio_generator = target_second_signal_generator(
-        audio_dir.as_posix(),
-        duration=args.duration,
-        sample_rate=args.target_sample_rate,
-        max_epoch=1,
-    )
-    count = 0
-    process_bar = tqdm(desc="build dataset jsonl")
-    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
-        for audio in audio_generator:
-            if count >= args.max_count > 0:
-                break
-
-            filename = audio["filename"]
-            raw_duration = audio["raw_duration"]
-            offset = audio["offset"]
-            duration = audio["duration"]
-
-            random1 = random.random()
-            random2 = random.random()
-
-            row = {
-                "count": count,
-
-                "filename": filename,
-                "raw_duration": raw_duration,
-                "offset": offset,
-                "duration": duration,
-
-                "random1": random1,
-            }
-            row = json.dumps(row, ensure_ascii=False)
-            if random2 < (1 / 10):
-                fvalid.write(f"{row}\n")
-            else:
-                ftrain.write(f"{row}\n")
-
-            count += 1
-            duration_seconds = count * args.duration
-            duration_hours = duration_seconds / 3600
-
-            process_bar.update(n=1)
-            process_bar.set_postfix({
-                "duration_hours": round(duration_hours, 4),
-            })
-
-    return
-
-
-if __name__ == "__main__":
-    main()

examples/frcrn_mp3_to_wav/step_2_train_model.py

@@ -1,442 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import argparse
-import json
-import logging
-from logging.handlers import TimedRotatingFileHandler
-import os
-import platform
-from pathlib import Path
-import random
-import sys
-import shutil
-from typing import List
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-from torch.utils.data.dataloader import DataLoader
-from tqdm import tqdm
-
-from toolbox.torch.utils.data.dataset.mp3_to_wav_jsonl_dataset import Mp3ToWavJsonlDataset
-from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
-from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
-from toolbox.torchaudio.metrics.pesq import run_pesq_score
-from toolbox.torchaudio.models.frcrn.configuration_frcrn import FRCRNConfig
-from toolbox.torchaudio.models.frcrn.modeling_frcrn import FRCRN, FRCRNPretrainedModel
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
-    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
-
-    parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
-    parser.add_argument("--patience", default=30, type=int)
-    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
-
-    parser.add_argument("--config_file", default="config.yaml", type=str)
-
-    args = parser.parse_args()
-    return args
-
-
-def logging_config(file_dir: str):
-    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
-
-    logging.basicConfig(format=fmt,
-                        datefmt="%m/%d/%Y %H:%M:%S",
-                        level=logging.INFO)
-    file_handler = TimedRotatingFileHandler(
-        filename=os.path.join(file_dir, "main.log"),
-        encoding="utf-8",
-        when="D",
-        interval=1,
-        backupCount=7
-    )
-    file_handler.setLevel(logging.INFO)
-    file_handler.setFormatter(logging.Formatter(fmt))
-    logger = logging.getLogger(__name__)
-    logger.addHandler(file_handler)
-
-    return logger
-
-
-class CollateFunction(object):
-    def __init__(self):
-        pass
-
-    def __call__(self, batch: List[dict]):
-        mp3_waveform_list = list()
-        wav_waveform_list = list()
-
-        for sample in batch:
-            mp3_waveform: torch.Tensor = sample["mp3_waveform"]
-            wav_waveform: torch.Tensor = sample["wav_waveform"]
-
-            mp3_waveform_list.append(mp3_waveform)
-            wav_waveform_list.append(wav_waveform)
-
-        mp3_waveform_list = torch.stack(mp3_waveform_list)
-        wav_waveform_list = torch.stack(wav_waveform_list)
-
-        # assert
-        if torch.any(torch.isnan(mp3_waveform_list)) or torch.any(torch.isinf(mp3_waveform_list)):
-            raise AssertionError("nan or inf in mp3_waveform_list")
-        if torch.any(torch.isnan(wav_waveform_list)) or torch.any(torch.isinf(wav_waveform_list)):
-            raise AssertionError("nan or inf in wav_waveform_list")
-
-        return mp3_waveform_list, wav_waveform_list
-
-
-collate_fn = CollateFunction()
-
-
-def main():
-    args = get_args()
-
-    config = FRCRNConfig.from_pretrained(
-        pretrained_model_name_or_path=args.config_file,
-    )
-
-    serialization_dir = Path(args.serialization_dir)
-    serialization_dir.mkdir(parents=True, exist_ok=True)
-
-    logger = logging_config(serialization_dir)
-
-    random.seed(config.seed)
-    np.random.seed(config.seed)
-    torch.manual_seed(config.seed)
-    logger.info(f"set seed: {config.seed}")
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    n_gpu = torch.cuda.device_count()
-    logger.info(f"GPU available count: {n_gpu}; device: {device}")
-
-    # datasets
-    train_dataset = Mp3ToWavJsonlDataset(
-        jsonl_file=args.train_dataset,
-        expected_sample_rate=config.sample_rate,
-        max_wave_value=32768.0,
-        # skip=225000,
-    )
-    valid_dataset = Mp3ToWavJsonlDataset(
-        jsonl_file=args.valid_dataset,
-        expected_sample_rate=config.sample_rate,
-        max_wave_value=32768.0,
-    )
-    train_data_loader = DataLoader(
-        dataset=train_dataset,
-        batch_size=config.batch_size,
-        # shuffle=True,
-        sampler=None,
-        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
-        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
-        collate_fn=collate_fn,
-        pin_memory=False,
-        prefetch_factor=2,
-    )
-    valid_data_loader = DataLoader(
-        dataset=valid_dataset,
-        batch_size=config.batch_size,
-        # shuffle=True,
-        sampler=None,
-        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
-        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
-        collate_fn=collate_fn,
-        pin_memory=False,
-        prefetch_factor=2,
-    )
-
-    # models
-    logger.info(f"prepare models. config_file: {args.config_file}")
-    model = FRCRNPretrainedModel(config).to(device)
-    model.to(device)
-    model.train()
-
-    # optimizer
-    logger.info("prepare optimizer, lr_scheduler, loss_fn, evaluation_metric")
-    optimizer = torch.optim.AdamW(model.get_params(weight_decay=config.weight_decay), config.lr)
-
-    # resume training
-    last_step_idx = -1
-    last_epoch = -1
-    for step_idx_str in serialization_dir.glob("steps-*"):
-        step_idx_str = Path(step_idx_str)
-        step_idx = step_idx_str.stem.split("-")[1]
-        step_idx = int(step_idx)
-        if step_idx > last_step_idx:
-            last_step_idx = step_idx
-    # last_epoch = 0
-
-    if last_step_idx != -1:
-        logger.info(f"resume from steps-{last_step_idx}.")
-        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
-        # optimizer_pth = serialization_dir / f"steps-{last_step_idx}/optimizer.pth"
-
-        logger.info(f"load state dict for model.")
-        with open(model_pt.as_posix(), "rb") as f:
-            state_dict = torch.load(f, map_location="cpu", weights_only=True)
-        model.load_state_dict(state_dict, strict=True)
-
-        # logger.info(f"load state dict for optimizer.")
-        # with open(optimizer_pth.as_posix(), "rb") as f:
-        #     state_dict = torch.load(f, map_location="cpu", weights_only=True)
-        # optimizer.load_state_dict(state_dict)
-
-    if config.lr_scheduler == "CosineAnnealingLR":
-        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
-            optimizer,
-            last_epoch=last_epoch,
-            # T_max=10 * config.eval_steps,
-            # eta_min=0.01 * config.lr,
-            **config.lr_scheduler_kwargs,
-        )
-    elif config.lr_scheduler == "MultiStepLR":
-        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-            optimizer,
-            last_epoch=last_epoch,
-            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
-        )
-    else:
-        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
-
-    neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
-    mr_stft_loss_fn = MultiResolutionSTFTLoss(
-        fft_size_list=[256, 512, 1024],
-        win_size_list=[256, 512, 1024],
-        hop_size_list=[128, 256, 512],
-        factor_sc=1.5,
-        factor_mag=1.0,
-        reduction="mean"
-    ).to(device)
-
-    # training loop
-
-    # state
-    average_pesq_score = 1000000000
-    average_loss = 1000000000
-    average_neg_si_snr_loss = 1000000000
-    average_mask_loss = 1000000000
-
-    model_list = list()
-    best_epoch_idx = None
-    best_step_idx = None
-    best_metric = None
-    patience_count = 0
-
-    step_idx = 0 if last_step_idx == -1 else last_step_idx
-
-    logger.info("training")
-    early_stop_flag = False
-    for epoch_idx in range(max(0, last_epoch+1), config.max_epochs):
-        if early_stop_flag:
-            break
-
-        # train
-        model.train()
-
-        total_pesq_score = 0.
-        total_loss = 0.
-        total_mr_stft_loss = 0.
-        total_neg_si_snr_loss = 0.
-        total_mask_loss = 0.
-        total_batches = 0.
-
-        progress_bar_train = tqdm(
-            initial=step_idx,
-            desc="Training; epoch-{}".format(epoch_idx),
-        )
-        for train_batch in train_data_loader:
-            mp3_audios, wav_audios = train_batch
-            noisy_audios: torch.Tensor = mp3_audios.to(device)
-            clean_audios: torch.Tensor = wav_audios.to(device)
-
-            est_spec, est_wav, est_mask = model.forward(noisy_audios)
-            denoise_audios = est_wav
-
-            mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
-            neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
-            mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
-
-            loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss
-            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
-                logger.info(f"find nan or inf in loss.")
-                continue
-
-            denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
-            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
-
-            optimizer.zero_grad()
-            loss.backward()
-            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.clip_grad_norm)
-            optimizer.step()
-            lr_scheduler.step()
-
-            total_pesq_score += pesq_score
-            total_loss += loss.item()
-            total_mr_stft_loss += mr_stft_loss.item()
-            total_neg_si_snr_loss += neg_si_snr_loss.item()
-            total_mask_loss += mask_loss.item()
-            total_batches += 1
-
-            average_pesq_score = round(total_pesq_score / total_batches, 4)
-            average_loss = round(total_loss / total_batches, 4)
-            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
-            average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
-            average_mask_loss = round(total_mask_loss / total_batches, 4)
-
-            progress_bar_train.update(1)
-            progress_bar_train.set_postfix({
-                "lr": lr_scheduler.get_last_lr()[0],
-                "pesq_score": average_pesq_score,
-                "loss": average_loss,
-                "mr_stft_loss": average_mr_stft_loss,
-                "neg_si_snr_loss": average_neg_si_snr_loss,
-                "mask_loss": average_mask_loss,
-            })
-
-            # evaluation
-            step_idx += 1
-            if step_idx % config.eval_steps == 0:
-                model.eval()
-                with torch.no_grad():
-                    torch.cuda.empty_cache()
-
-                    total_pesq_score = 0.
-                    total_loss = 0.
-                    total_mr_stft_loss = 0.
-                    total_neg_si_snr_loss = 0.
-                    total_mask_loss = 0.
-                    total_batches = 0.
-
-                    progress_bar_train.close()
-                    progress_bar_eval = tqdm(
-                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
-                    )
-                    for eval_batch in valid_data_loader:
-                        mp3_audios, wav_audios = eval_batch
-                        noisy_audios: torch.Tensor = mp3_audios.to(device)
-                        clean_audios: torch.Tensor = wav_audios.to(device)
-
-                        est_spec, est_wav, est_mask = model.forward(noisy_audios)
-                        denoise_audios = est_wav
-
-                        mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
-                        neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
-                        mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
-
-                        loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss
-                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
-                            logger.info(f"find nan or inf in loss.")
-                            continue
-
-                        denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
-                        clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
-
-                        total_pesq_score += pesq_score
-                        total_loss += loss.item()
-                        total_neg_si_snr_loss += neg_si_snr_loss.item()
-                        total_mask_loss += mask_loss.item()
-                        total_batches += 1
-
-                        average_pesq_score = round(total_pesq_score / total_batches, 4)
-                        average_loss = round(total_loss / total_batches, 4)
-                        average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
-                        average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
-                        average_mask_loss = round(total_mask_loss / total_batches, 4)
-
-                        progress_bar_eval.update(1)
-                        progress_bar_eval.set_postfix({
-                            "lr": lr_scheduler.get_last_lr()[0],
-                            "pesq_score": average_pesq_score,
-                            "loss": average_loss,
-                            "mr_stft_loss": average_mr_stft_loss,
-                            "neg_si_snr_loss": average_neg_si_snr_loss,
-                            "mask_loss": average_mask_loss,
-                        })
-
-                    total_pesq_score = 0.
-                    total_loss = 0.
-                    total_mr_stft_loss = 0.
-                    total_neg_si_snr_loss = 0.
-                    total_mask_loss = 0.
-                    total_batches = 0.
-
-                    progress_bar_eval.close()
-                    progress_bar_train = tqdm(
-                        initial=progress_bar_train.n,
-                        postfix=progress_bar_train.postfix,
-                        desc=progress_bar_train.desc,
-                    )
-
-                    # save path
-                    save_dir = serialization_dir / "steps-{}".format(step_idx)
-                    save_dir.mkdir(parents=True, exist_ok=False)
-
-                    # save models
-                    model.save_pretrained(save_dir.as_posix())
-
-                    model_list.append(save_dir)
-                    if len(model_list) >= args.num_serialized_models_to_keep:
-                        model_to_delete: Path = model_list.pop(0)
-                        shutil.rmtree(model_to_delete.as_posix())
-
-                    # save metric
-                    if best_metric is None:
-                        best_epoch_idx = epoch_idx
-                        best_step_idx = step_idx
-                        best_metric = average_pesq_score
-                    elif average_pesq_score >= best_metric:
-                        # great is better.
-                        best_epoch_idx = epoch_idx
-                        best_step_idx = step_idx
-                        best_metric = average_pesq_score
-                    else:
-                        pass
-
-                    metrics = {
-                        "epoch_idx": epoch_idx,
-                        "best_epoch_idx": best_epoch_idx,
-                        "best_step_idx": best_step_idx,
-                        "pesq_score": average_pesq_score,
-                        "loss": average_loss,
-                        "neg_si_snr_loss": average_neg_si_snr_loss,
-                        "mask_loss": average_mask_loss,
-                    }
-                    metrics_filename = save_dir / "metrics_epoch.json"
-                    with open(metrics_filename, "w", encoding="utf-8") as f:
-                        json.dump(metrics, f, indent=4, ensure_ascii=False)
-
-                    # save best
-                    best_dir = serialization_dir / "best"
-                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
-                        if best_dir.exists():
-                            shutil.rmtree(best_dir)
-                        shutil.copytree(save_dir, best_dir)
-
-                    # early stop
-                    early_stop_flag = False
-                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
-                        patience_count = 0
-                    else:
-                        patience_count += 1
-                    if patience_count >= args.patience:
-                        early_stop_flag = True
-
-                    # early stop
-                    if early_stop_flag:
-                        break
-                model.train()
-
-    return
-
-
-if __name__ == "__main__":
-    main()

examples/frcrn_mp3_to_wav/yaml/config-10.yaml

@@ -1,31 +0,0 @@
-model_name: "frcrn"
-
-sample_rate: 8000
-segment_size: 32000
-nfft: 128
-win_size: 128
-hop_size: 64
-win_type: hann
-
-use_complex_networks: true
-model_depth: 10
-model_complexity: -1
-
-min_snr_db: -10
-max_snr_db: 20
-
-num_workers: 8
-batch_size: 32
-eval_steps: 20000
-
-lr: 0.001
-lr_scheduler: "CosineAnnealingLR"
-lr_scheduler_kwargs:
-  T_max: 250000
-  eta_min: 0.0001
-
-max_epochs: 100
-weight_decay: 1.0e-05
-clip_grad_norm: 10.0
-seed: 1234
-num_gpus: -1

examples/frcrn_mp3_to_wav/yaml/config-14.yaml

@@ -1,31 +0,0 @@
-model_name: "frcrn"
-
-sample_rate: 8000
-segment_size: 32000
-nfft: 640
-win_size: 640
-hop_size: 320
-win_type: hann
-
-use_complex_networks: true
-model_depth: 14
-model_complexity: -1
-
-min_snr_db: -10
-max_snr_db: 20
-
-num_workers: 8
-batch_size: 32
-eval_steps: 10000
-
-lr: 0.001
-lr_scheduler: "CosineAnnealingLR"
-lr_scheduler_kwargs:
-  T_max: 250000
-  eta_min: 0.0001
-
-max_epochs: 100
-weight_decay: 1.0e-05
-clip_grad_norm: 10.0
-seed: 1234
-num_gpus: -1

examples/frcrn_mp3_to_wav/yaml/config-20.yaml

@@ -1,31 +0,0 @@
-model_name: "frcrn"
-
-sample_rate: 8000
-segment_size: 32000
-nfft: 512
-win_size: 512
-hop_size: 256
-win_type: hann
-
-use_complex_networks: true
-model_depth: 20
-model_complexity: 45
-
-min_snr_db: -10
-max_snr_db: 20
-
-num_workers: 8
-batch_size: 32
-eval_steps: 10000
-
-lr: 0.001
-lr_scheduler: "CosineAnnealingLR"
-lr_scheduler_kwargs:
-  T_max: 250000
-  eta_min: 0.0001
-
-max_epochs: 100
-weight_decay: 1.0e-05
-clip_grad_norm: 10.0
-seed: 1234
-num_gpus: -1

examples/simple_linear_irm_aishell/run.sh

@@ -1,172 +0,0 @@
-#!/usr/bin/env bash
-
-: <<'END'
-
-sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name file_dir
-
-sh run.sh --stage 3 --stop_stage 3 --system_version windows --file_folder_name file_dir
-
-sh run.sh --stage 1 --stop_stage 3 --system_version centos --file_folder_name file_dir \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
-
-
-END
-
-
-# params
-system_version="windows";
-verbose=true;
-stage=0 # start from 0 if you need to start from data preparation
-stop_stage=9
-
-work_dir="$(pwd)"
-file_folder_name=file_folder_name
-final_model_name=final_model_name
-config_file="yaml/config.yaml"
-limit=10
-
-noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
-speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
-
-nohup_name=nohup.out
-
-# model params
-batch_size=64
-max_epochs=200
-save_top_k=10
-patience=5
-
-
-# parse options
-while true; do
-  [ -z "${1:-}" ] && break;  # break if there are no arguments
-  case "$1" in
-    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
-      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
-      old_value="(eval echo \\$$name)";
-      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
-        was_bool=true;
-      else
-        was_bool=false;
-      fi
-
-      # Set the variable to the right value-- the escaped quotes make it work if
-      # the option had spaces, like --cmd "queue.pl -sync y"
-      eval "${name}=\"$2\"";
-
-      # Check that Boolean-valued arguments are really Boolean.
-      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
-        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
-        exit 1;
-      fi
-      shift 2;
-      ;;
-
-    *) break;
-  esac
-done
-
-file_dir="${work_dir}/${file_folder_name}"
-final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
-evaluation_audio_dir="${file_dir}/evaluation_audio"
-
-dataset="${file_dir}/dataset.xlsx"
-train_dataset="${file_dir}/train.xlsx"
-valid_dataset="${file_dir}/valid.xlsx"
-
-$verbose && echo "system_version: ${system_version}"
-$verbose && echo "file_folder_name: ${file_folder_name}"
-
-if [ $system_version == "windows" ]; then
-  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
-elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
-  #source /data/local/bin/nx_denoise/bin/activate
-  alias python3='/data/local/bin/nx_denoise/bin/python3'
-fi
-
-
-if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
-  $verbose && echo "stage 1: prepare data"
-  cd "${work_dir}" || exit 1
-  python3 step_1_prepare_data.py \
-  --file_dir "${file_dir}" \
-  --noise_dir "${noise_dir}" \
-  --speech_dir "${speech_dir}" \
-  --train_dataset "${train_dataset}" \
-  --valid_dataset "${valid_dataset}" \
-
-fi
-
-
-if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
-  $verbose && echo "stage 2: train model"
-  cd "${work_dir}" || exit 1
-  python3 step_2_train_model.py \
-  --train_dataset "${train_dataset}" \
-  --valid_dataset "${valid_dataset}" \
-  --serialization_dir "${file_dir}" \
-  --config_file "${config_file}" \
-
-fi
-
-
-if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
-  $verbose && echo "stage 3: test model"
-  cd "${work_dir}" || exit 1
-  python3 step_3_evaluation.py \
-  --valid_dataset "${valid_dataset}" \
-  --model_dir "${file_dir}/best" \
-  --evaluation_audio_dir "${evaluation_audio_dir}" \
-  --limit "${limit}" \
-
-fi
-
-
-if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
-  $verbose && echo "stage 4: export model"
-  cd "${work_dir}" || exit 1
-  python3 step_5_export_models.py \
-  --vocabulary_dir "${vocabulary_dir}" \
-  --model_dir "${file_dir}/best" \
-  --serialization_dir "${file_dir}" \
-
-fi
-
-
-if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
-  $verbose && echo "stage 5: collect files"
-  cd "${work_dir}" || exit 1
-
-  mkdir -p ${final_model_dir}
-
-  cp "${file_dir}/best"/* "${final_model_dir}"
-  cp -r "${file_dir}/vocabulary" "${final_model_dir}"
-
-  cp "${file_dir}/evaluation.xlsx" "${final_model_dir}/evaluation.xlsx"
-
-  cp "${file_dir}/trace_model.zip" "${final_model_dir}/trace_model.zip"
-  cp "${file_dir}/trace_quant_model.zip" "${final_model_dir}/trace_quant_model.zip"
-  cp "${file_dir}/script_model.zip" "${final_model_dir}/script_model.zip"
-  cp "${file_dir}/script_quant_model.zip" "${final_model_dir}/script_quant_model.zip"
-
-  cd "${final_model_dir}/.." || exit 1;
-
-  if [ -e "${final_model_name}.zip" ]; then
-    rm -rf "${final_model_name}_backup.zip"
-    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
-  fi
-
-  zip -r "${final_model_name}.zip" "${final_model_name}"
-  rm -rf "${final_model_name}"
-
-fi
-
-
-if [ ${stage} -le 6 ] && [ ${stop_stage} -ge 6 ]; then
-  $verbose && echo "stage 6: clear file_dir"
-  cd "${work_dir}" || exit 1
-
-  rm -rf "${file_dir}";
-
-fi

examples/simple_linear_irm_aishell/step_1_prepare_data.py

@@ -1,196 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import argparse
-import os
-from pathlib import Path
-import random
-import sys
-import shutil
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import pandas as pd
-from scipy.io import wavfile
-from tqdm import tqdm
-import librosa
-
-from project_settings import project_path
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--file_dir", default="./", type=str)
-
-    parser.add_argument(
-        "--noise_dir",
-        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
-        type=str
-    )
-    parser.add_argument(
-        "--speech_dir",
-        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
-        type=str
-    )
-
-    parser.add_argument("--train_dataset", default="train.xlsx", type=str)
-    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
-
-    parser.add_argument("--duration", default=2.0, type=float)
-    parser.add_argument("--min_nsr_db", default=-20, type=float)
-    parser.add_argument("--max_nsr_db", default=5, type=float)
-
-    parser.add_argument("--target_sample_rate", default=8000, type=int)
-
-    args = parser.parse_args()
-    return args
-
-
-def filename_generator(data_dir: str):
-    data_dir = Path(data_dir)
-    for filename in data_dir.glob("**/*.wav"):
-        yield filename.as_posix()
-
-
-def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000):
-    data_dir = Path(data_dir)
-    for filename in data_dir.glob("**/*.wav"):
-        signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
-        raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
-
-        if raw_duration < duration:
-            # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
-            continue
-        if signal.ndim != 1:
-            raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
-
-        signal_length = len(signal)
-        win_size = int(duration * sample_rate)
-        for begin in range(0, signal_length - win_size, win_size):
-            row = {
-                "filename": filename.as_posix(),
-                "raw_duration": round(raw_duration, 4),
-                "offset": round(begin / sample_rate, 4),
-                "duration": round(duration, 4),
-            }
-            yield row
-
-
-def get_dataset(args):
-    file_dir = Path(args.file_dir)
-    file_dir.mkdir(exist_ok=True)
-
-    noise_dir = Path(args.noise_dir)
-    speech_dir = Path(args.speech_dir)
-
-    noise_generator = target_second_signal_generator(
-        noise_dir.as_posix(),
-        duration=args.duration,
-        sample_rate=args.target_sample_rate
-    )
-    speech_generator = target_second_signal_generator(
-        speech_dir.as_posix(),
-        duration=args.duration,
-        sample_rate=args.target_sample_rate
-    )
-
-    dataset = list()
-
-    count = 0
-    process_bar = tqdm(desc="build dataset excel")
-    for noise, speech in zip(noise_generator, speech_generator):
-
-        noise_filename = noise["filename"]
-        noise_raw_duration = noise["raw_duration"]
-        noise_offset = noise["offset"]
-        noise_duration = noise["duration"]
-
-        speech_filename = speech["filename"]
-        speech_raw_duration = speech["raw_duration"]
-        speech_offset = speech["offset"]
-        speech_duration = speech["duration"]
-
-        random1 = random.random()
-        random2 = random.random()
-
-        row = {
-            "noise_filename": noise_filename,
-            "noise_raw_duration": noise_raw_duration,
-            "noise_offset": noise_offset,
-            "noise_duration": noise_duration,
-
-            "speech_filename": speech_filename,
-            "speech_raw_duration": speech_raw_duration,
-            "speech_offset": speech_offset,
-            "speech_duration": speech_duration,
-
-            "snr_db": random.uniform(args.min_nsr_db, args.max_nsr_db),
-
-            "random1": random1,
-            "random2": random2,
-            "flag": "TRAIN" if random2 < 0.8 else "TEST",
-        }
-        dataset.append(row)
-        count += 1
-        duration_seconds = count * args.duration
-        duration_hours = duration_seconds / 3600
-
-        process_bar.update(n=1)
-        process_bar.set_postfix({
-            # "duration_seconds": round(duration_seconds, 4),
-            "duration_hours": round(duration_hours, 4),
-        })
-
-    dataset = pd.DataFrame(dataset)
-    dataset = dataset.sort_values(by=["random1"], ascending=False)
-    dataset.to_excel(
-        file_dir / "dataset.xlsx",
-        index=False,
-    )
-    return
-
-
-
-def split_dataset(args):
-    """分割训练集, 测试集"""
-    file_dir = Path(args.file_dir)
-    file_dir.mkdir(exist_ok=True)
-
-    df = pd.read_excel(file_dir / "dataset.xlsx")
-
-    train = list()
-    test = list()
-
-    for i, row in df.iterrows():
-        flag = row["flag"]
-        if flag == "TRAIN":
-            train.append(row)
-        else:
-            test.append(row)
-
-    train = pd.DataFrame(train)
-    train.to_excel(
-        args.train_dataset,
-        index=False,
-        # encoding="utf_8_sig"
-    )
-    test = pd.DataFrame(test)
-    test.to_excel(
-        args.valid_dataset,
-        index=False,
-        # encoding="utf_8_sig"
-    )
-
-    return
-
-
-def main():
-    args = get_args()
-
-    get_dataset(args)
-    split_dataset(args)
-    return
-
-
-if __name__ == "__main__":
-    main()

examples/simple_linear_irm_aishell/step_2_train_model.py

@@ -1,348 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-"""
-https://github.com/WenzheLiu-Speech/awesome-speech-enhancement
-"""
-import argparse
-import json
-import logging
-from logging.handlers import TimedRotatingFileHandler
-import os
-import platform
-from pathlib import Path
-import random
-import sys
-import shutil
-from typing import List
-
-from torch import dtype
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.utils.data.dataloader import DataLoader
-import torchaudio
-from tqdm import tqdm
-
-from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
-from toolbox.torchaudio.models.simple_linear_irm.configuration_simple_linear_irm import SimpleLinearIRMConfig
-from toolbox.torchaudio.models.simple_linear_irm.modeling_simple_linear_irm import SimpleLinearIRMPretrainedModel
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--train_dataset", default="train.xlsx", type=str)
-    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
-
-    parser.add_argument("--max_epochs", default=100, type=int)
-
-    parser.add_argument("--batch_size", default=64, type=int)
-    parser.add_argument("--learning_rate", default=1e-3, type=float)
-    parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
-    parser.add_argument("--patience", default=5, type=int)
-    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
-    parser.add_argument("--seed", default=0, type=int)
-
-    parser.add_argument("--config_file", default="config.yaml", type=str)
-
-    args = parser.parse_args()
-    return args
-
-
-def logging_config(file_dir: str):
-    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
-
-    logging.basicConfig(format=fmt,
-                        datefmt="%m/%d/%Y %H:%M:%S",
-                        level=logging.INFO)
-    file_handler = TimedRotatingFileHandler(
-        filename=os.path.join(file_dir, "main.log"),
-        encoding="utf-8",
-        when="D",
-        interval=1,
-        backupCount=7
-    )
-    file_handler.setLevel(logging.INFO)
-    file_handler.setFormatter(logging.Formatter(fmt))
-    logger = logging.getLogger(__name__)
-    logger.addHandler(file_handler)
-
-    return logger
-
-
-class CollateFunction(object):
-    def __init__(self,
-                 n_fft: int = 512,
-                 win_length: int = 200,
-                 hop_length: int = 80,
-                 window_fn: str = "hamming",
-                 irm_beta: float = 1.0,
-                 epsilon: float = 1e-8,
-                 ):
-        self.n_fft = n_fft
-        self.win_length = win_length
-        self.hop_length = hop_length
-        self.window_fn = window_fn
-        self.irm_beta = irm_beta
-        self.epsilon = epsilon
-
-        self.transform = torchaudio.transforms.Spectrogram(
-            n_fft=self.n_fft,
-            win_length=self.win_length,
-            hop_length=self.hop_length,
-            power=2.0,
-            window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
-        )
-
-    def __call__(self, batch: List[dict]):
-        mix_spec_list = list()
-        speech_irm_list = list()
-        snr_db_list = list()
-        for sample in batch:
-            noise_wave: torch.Tensor = sample["noise_wave"]
-            speech_wave: torch.Tensor = sample["speech_wave"]
-            mix_wave: torch.Tensor = sample["mix_wave"]
-            snr_db: float = sample["snr_db"]
-
-            noise_spec = self.transform.forward(noise_wave)
-            speech_spec = self.transform.forward(speech_wave)
-            mix_spec = self.transform.forward(mix_wave)
-
-            # noise_irm = noise_spec / (noise_spec + speech_spec)
-            speech_irm = speech_spec / (noise_spec + speech_spec + self.epsilon)
-            speech_irm = torch.pow(speech_irm, self.irm_beta)
-
-            mix_spec_list.append(mix_spec)
-            speech_irm_list.append(speech_irm)
-            snr_db_list.append(torch.tensor(snr_db, dtype=torch.float32))
-
-        mix_spec_list = torch.stack(mix_spec_list)
-        speech_irm_list = torch.stack(speech_irm_list)
-        snr_db_list = torch.stack(snr_db_list)  # shape: (batch_size,)
-
-        # assert
-        if torch.any(torch.isnan(mix_spec_list)):
-            raise AssertionError("nan in mix_spec Tensor")
-        if torch.any(torch.isnan(speech_irm_list)):
-            raise AssertionError("nan in speech_irm Tensor")
-        if torch.any(torch.isnan(snr_db_list)):
-            raise AssertionError("nan in snr_db Tensor")
-
-        return mix_spec_list, speech_irm_list, snr_db_list
-
-
-collate_fn = CollateFunction()
-
-
-def main():
-    args = get_args()
-
-    serialization_dir = Path(args.serialization_dir)
-    serialization_dir.mkdir(parents=True, exist_ok=True)
-
-    logger = logging_config(serialization_dir)
-
-    random.seed(args.seed)
-    np.random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    logger.info("set seed: {}".format(args.seed))
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    n_gpu = torch.cuda.device_count()
-    logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
-
-    # datasets
-    logger.info("prepare datasets")
-    train_dataset = DenoiseExcelDataset(
-        excel_file=args.train_dataset,
-        expected_sample_rate=8000,
-        max_wave_value=32768.0,
-    )
-    valid_dataset = DenoiseExcelDataset(
-        excel_file=args.valid_dataset,
-        expected_sample_rate=8000,
-        max_wave_value=32768.0,
-    )
-    train_data_loader = DataLoader(
-        dataset=train_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
-        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
-        collate_fn=collate_fn,
-        pin_memory=False,
-        # prefetch_factor=64,
-    )
-    valid_data_loader = DataLoader(
-        dataset=valid_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
-        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
-        collate_fn=collate_fn,
-        pin_memory=False,
-        # prefetch_factor=64,
-    )
-
-    # models
-    logger.info(f"prepare models. config_file: {args.config_file}")
-    config = SimpleLinearIRMConfig.from_pretrained(
-        pretrained_model_name_or_path=args.config_file,
-        # num_labels=vocabulary.get_vocab_size(namespace="labels")
-    )
-    model = SimpleLinearIRMPretrainedModel(
-        config=config,
-    )
-    model.to(device)
-    model.train()
-
-    # optimizer
-    logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
-    param_optimizer = model.parameters()
-    optimizer = torch.optim.Adam(
-        param_optimizer,
-        lr=args.learning_rate,
-    )
-    # lr_scheduler = torch.optim.lr_scheduler.StepLR(
-    #     optimizer,
-    #     step_size=2000
-    # )
-    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-        optimizer,
-        milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
-    )
-    mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-
-    # training loop
-    logger.info("training")
-
-    training_loss = 10000000000
-    evaluation_loss = 10000000000
-
-    model_list = list()
-    best_idx_epoch = None
-    best_metric = None
-    patience_count = 0
-
-    for idx_epoch in range(args.max_epochs):
-        total_loss = 0.
-        total_examples = 0.
-        progress_bar = tqdm(
-            total=len(train_data_loader),
-            desc="Training; epoch: {}".format(idx_epoch),
-        )
-
-        for batch in train_data_loader:
-            mix_spec, speech_irm, snr_db = batch
-            mix_spec = mix_spec.to(device)
-            speech_irm_target = speech_irm.to(device)
-            snr_db_target = snr_db.to(device)
-
-            speech_irm_prediction = model.forward(mix_spec)
-            loss = mse_loss.forward(speech_irm_prediction, speech_irm_target)
-
-            total_loss += loss.item()
-            total_examples += mix_spec.size(0)
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-            lr_scheduler.step()
-
-            training_loss = total_loss / total_examples
-            training_loss = round(training_loss, 4)
-
-            progress_bar.update(1)
-            progress_bar.set_postfix({
-                "training_loss": training_loss,
-            })
-
-        total_loss = 0.
-        total_examples = 0.
-        progress_bar = tqdm(
-            total=len(valid_data_loader),
-            desc="Evaluation; epoch: {}".format(idx_epoch),
-        )
-        for batch in valid_data_loader:
-            mix_spec, speech_irm, snr_db = batch
-            mix_spec = mix_spec.to(device)
-            speech_irm_target = speech_irm.to(device)
-            snr_db_target = snr_db.to(device)
-
-            with torch.no_grad():
-                speech_irm_prediction = model.forward(mix_spec)
-                loss = mse_loss.forward(speech_irm_prediction, speech_irm_target)
-
-            total_loss += loss.item()
-            total_examples += mix_spec.size(0)
-
-            evaluation_loss = total_loss / total_examples
-            evaluation_loss = round(evaluation_loss, 4)
-
-            progress_bar.update(1)
-            progress_bar.set_postfix({
-                "evaluation_loss": evaluation_loss,
-            })
-
-        # save path
-        epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
-        epoch_dir.mkdir(parents=True, exist_ok=False)
-
-        # save models
-        model.save_pretrained(epoch_dir.as_posix())
-
-        model_list.append(epoch_dir)
-        if len(model_list) >= args.num_serialized_models_to_keep:
-            model_to_delete: Path = model_list.pop(0)
-            shutil.rmtree(model_to_delete.as_posix())
-
-        # save metric
-        if best_metric is None:
-            best_idx_epoch = idx_epoch
-            best_metric = evaluation_loss
-        elif evaluation_loss < best_metric:
-            best_idx_epoch = idx_epoch
-            best_metric = evaluation_loss
-        else:
-            pass
-
-        metrics = {
-            "idx_epoch": idx_epoch,
-            "best_idx_epoch": best_idx_epoch,
-            "training_loss": training_loss,
-            "evaluation_loss": evaluation_loss,
-            "learning_rate": optimizer.param_groups[0]["lr"],
-        }
-        metrics_filename = epoch_dir / "metrics_epoch.json"
-        with open(metrics_filename, "w", encoding="utf-8") as f:
-            json.dump(metrics, f, indent=4, ensure_ascii=False)
-
-        # save best
-        best_dir = serialization_dir / "best"
-        if best_idx_epoch == idx_epoch:
-            if best_dir.exists():
-                shutil.rmtree(best_dir)
-            shutil.copytree(epoch_dir, best_dir)
-
-        # early stop
-        early_stop_flag = False
-        if best_idx_epoch == idx_epoch:
-            patience_count = 0
-        else:
-            patience_count += 1
-        if patience_count >= args.patience:
-            early_stop_flag = True
-
-        # early stop
-        if early_stop_flag:
-            break
-    return
-
-
-if __name__ == '__main__':
-    main()

examples/simple_linear_irm_aishell/step_3_evaluation.py

@@ -1,239 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import argparse
-import logging
-import os
-from pathlib import Path
-import sys
-import uuid
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import librosa
-import numpy as np
-import pandas as pd
-from scipy.io import wavfile
-import torch
-import torch.nn as nn
-import torchaudio
-from tqdm import tqdm
-
-from toolbox.torchaudio.models.simple_linear_irm.modeling_simple_linear_irm import SimpleLinearIRMPretrainedModel
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
-    parser.add_argument("--model_dir", default="serialization_dir/best", type=str)
-    parser.add_argument("--evaluation_audio_dir", default="evaluation_audio_dir", type=str)
-
-    parser.add_argument("--limit", default=10, type=int)
-
-    args = parser.parse_args()
-    return args
-
-
-def logging_config():
-    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
-
-    logging.basicConfig(format=fmt,
-                        datefmt="%m/%d/%Y %H:%M:%S",
-                        level=logging.INFO)
-    stream_handler = logging.StreamHandler()
-    stream_handler.setLevel(logging.INFO)
-    stream_handler.setFormatter(logging.Formatter(fmt))
-
-    logger = logging.getLogger(__name__)
-
-    return logger
-
-
-def mix_speech_and_noise(speech: np.ndarray, noise: np.ndarray, snr_db: float):
-    l1 = len(speech)
-    l2 = len(noise)
-    l = min(l1, l2)
-    speech = speech[:l]
-    noise = noise[:l]
-
-    # np.float32, value between (-1, 1).
-
-    speech_power = np.mean(np.square(speech))
-    noise_power = speech_power / (10 ** (snr_db / 10))
-
-    noise_adjusted = np.sqrt(noise_power) * noise / np.sqrt(np.mean(noise ** 2))
-
-    noisy_signal = speech + noise_adjusted
-
-    return noisy_signal
-
-
-stft_power = torchaudio.transforms.Spectrogram(
-    n_fft=512,
-    win_length=200,
-    hop_length=80,
-    power=2.0,
-    window_fn=torch.hamming_window,
-)
-
-
-stft_complex = torchaudio.transforms.Spectrogram(
-    n_fft=512,
-    win_length=200,
-    hop_length=80,
-    power=None,
-    window_fn=torch.hamming_window,
-)
-
-
-istft = torchaudio.transforms.InverseSpectrogram(
-    n_fft=512,
-    win_length=200,
-    hop_length=80,
-    window_fn=torch.hamming_window,
-)
-
-
-def enhance(mix_spec_complex: torch.Tensor, speech_irm_prediction: torch.Tensor):
-    mix_spec_complex = mix_spec_complex.detach().cpu()
-    speech_irm_prediction = speech_irm_prediction.detach().cpu()
-
-    mask_speech = speech_irm_prediction
-    mask_noise = 1.0 - speech_irm_prediction
-
-    speech_spec = mix_spec_complex * mask_speech
-    noise_spec = mix_spec_complex * mask_noise
-
-    speech_wave = istft.forward(speech_spec)
-    noise_wave = istft.forward(noise_spec)
-
-    return speech_wave, noise_wave
-
-
-def save_audios(noise_wave: torch.Tensor,
-                speech_wave: torch.Tensor,
-                mix_wave: torch.Tensor,
-                speech_wave_enhanced: torch.Tensor,
-                noise_wave_enhanced: torch.Tensor,
-                output_dir: str,
-                sample_rate: int = 8000,
-                ):
-    basename = uuid.uuid4().__str__()
-    output_dir = Path(output_dir) / basename
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    filename = output_dir / "noise_wave.wav"
-    torchaudio.save(filename, noise_wave, sample_rate)
-    filename = output_dir / "speech_wave.wav"
-    torchaudio.save(filename, speech_wave, sample_rate)
-    filename = output_dir / "mix_wave.wav"
-    torchaudio.save(filename, mix_wave, sample_rate)
-
-    filename = output_dir / "speech_wave_enhanced.wav"
-    torchaudio.save(filename, speech_wave_enhanced, sample_rate)
-    filename = output_dir / "noise_wave_enhanced.wav"
-    torchaudio.save(filename, noise_wave_enhanced, sample_rate)
-
-    return output_dir.as_posix()
-
-
-def main():
-    args = get_args()
-
-    logger = logging_config()
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    n_gpu = torch.cuda.device_count()
-    logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
-
-    logger.info("prepare model")
-    model = SimpleLinearIRMPretrainedModel.from_pretrained(
-        pretrained_model_name_or_path=args.model_dir,
-    )
-    model.to(device)
-    model.eval()
-
-    # optimizer
-    logger.info("prepare loss_fn")
-    mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-
-    logger.info("read excel")
-    df = pd.read_excel(args.valid_dataset)
-
-    total_loss = 0.
-    total_examples = 0.
-    progress_bar = tqdm(total=len(df), desc="Evaluation")
-    for idx, row in df.iterrows():
-        noise_filename = row["noise_filename"]
-        noise_offset = row["noise_offset"]
-        noise_duration = row["noise_duration"]
-
-        speech_filename = row["speech_filename"]
-        speech_offset = row["speech_offset"]
-        speech_duration = row["speech_duration"]
-
-        snr_db = row["snr_db"]
-
-        noise_wave, _ = librosa.load(
-            noise_filename,
-            sr=8000,
-            offset=noise_offset,
-            duration=noise_duration,
-        )
-        speech_wave, _ = librosa.load(
-            speech_filename,
-            sr=8000,
-            offset=speech_offset,
-            duration=speech_duration,
-        )
-        mix_wave: np.ndarray = mix_speech_and_noise(
-            speech=speech_wave,
-            noise=noise_wave,
-            snr_db=snr_db,
-        )
-        noise_wave = torch.tensor(noise_wave, dtype=torch.float32)
-        speech_wave = torch.tensor(speech_wave, dtype=torch.float32)
-        mix_wave: torch.Tensor = torch.tensor(mix_wave, dtype=torch.float32)
-
-        noise_wave = noise_wave.unsqueeze(dim=0)
-        speech_wave = speech_wave.unsqueeze(dim=0)
-        mix_wave = mix_wave.unsqueeze(dim=0)
-
-        noise_spec: torch.Tensor = stft_power.forward(noise_wave)
-        speech_spec: torch.Tensor = stft_power.forward(speech_wave)
-        mix_spec: torch.Tensor = stft_power.forward(mix_wave)
-        mix_spec_complex: torch.Tensor = stft_complex.forward(mix_wave)
-
-        speech_irm = speech_spec / (noise_spec + speech_spec)
-        speech_irm = torch.pow(speech_irm, 1.0)
-
-        mix_spec = mix_spec.to(device)
-        speech_irm_target = speech_irm.to(device)
-        with torch.no_grad():
-            speech_irm_prediction = model.forward(mix_spec)
-            loss = mse_loss.forward(speech_irm_prediction, speech_irm_target)
-
-        speech_wave_enhanced, noise_wave_enhanced = enhance(mix_spec_complex, speech_irm_prediction)
-        save_audios(noise_wave, speech_wave, mix_wave, speech_wave_enhanced, noise_wave_enhanced, args.evaluation_audio_dir)
-
-        total_loss += loss.item()
-        total_examples += mix_spec.size(0)
-
-        evaluation_loss = total_loss / total_examples
-        evaluation_loss = round(evaluation_loss, 4)
-
-        progress_bar.update(1)
-        progress_bar.set_postfix({
-            "evaluation_loss": evaluation_loss,
-        })
-
-        if idx > args.limit:
-            break
-
-    return
-
-
-if __name__ == '__main__':
-    main()

examples/simple_linear_irm_aishell/yaml/config.yaml

@@ -1,13 +0,0 @@
-model_name: "simple_linear_irm"
-
-# spec
-sample_rate: 8000
-n_fft: 512
-win_length: 200
-hop_length: 80
-
-# model
-num_bins: 257
-hidden_size: 2048
-lookback: 3
-lookahead: 3

examples/spectrum_dfnet_aishell/run.sh

@@ -1,178 +0,0 @@
-#!/usr/bin/env bash
-
-: <<'END'
-
-
-sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name file_dir \
---noise_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "E:/programmer/asr_datasets/aishell/data_aishell/wav/train"
-
-
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
-
-sh run.sh --stage 3 --stop_stage 3 --system_version centos --file_folder_name file_dir \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
-
-
-END
-
-
-# params
-system_version="windows";
-verbose=true;
-stage=0 # start from 0 if you need to start from data preparation
-stop_stage=9
-
-work_dir="$(pwd)"
-file_folder_name=file_folder_name
-final_model_name=final_model_name
-config_file="yaml/config.yaml"
-limit=10
-
-noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
-speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
-
-nohup_name=nohup.out
-
-# model params
-batch_size=64
-max_epochs=200
-save_top_k=10
-patience=5
-
-
-# parse options
-while true; do
-  [ -z "${1:-}" ] && break;  # break if there are no arguments
-  case "$1" in
-    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
-      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
-      old_value="(eval echo \\$$name)";
-      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
-        was_bool=true;
-      else
-        was_bool=false;
-      fi
-
-      # Set the variable to the right value-- the escaped quotes make it work if
-      # the option had spaces, like --cmd "queue.pl -sync y"
-      eval "${name}=\"$2\"";
-
-      # Check that Boolean-valued arguments are really Boolean.
-      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
-        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
-        exit 1;
-      fi
-      shift 2;
-      ;;
-
-    *) break;
-  esac
-done
-
-file_dir="${work_dir}/${file_folder_name}"
-final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
-evaluation_audio_dir="${file_dir}/evaluation_audio"
-
-dataset="${file_dir}/dataset.xlsx"
-train_dataset="${file_dir}/train.xlsx"
-valid_dataset="${file_dir}/valid.xlsx"
-
-$verbose && echo "system_version: ${system_version}"
-$verbose && echo "file_folder_name: ${file_folder_name}"
-
-if [ $system_version == "windows" ]; then
-  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
-elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
-  #source /data/local/bin/nx_denoise/bin/activate
-  alias python3='/data/local/bin/nx_denoise/bin/python3'
-fi
-
-
-if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
-  $verbose && echo "stage 1: prepare data"
-  cd "${work_dir}" || exit 1
-  python3 step_1_prepare_data.py \
-  --file_dir "${file_dir}" \
-  --noise_dir "${noise_dir}" \
-  --speech_dir "${speech_dir}" \
-  --train_dataset "${train_dataset}" \
-  --valid_dataset "${valid_dataset}" \
-
-fi
-
-
-if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
-  $verbose && echo "stage 2: train model"
-  cd "${work_dir}" || exit 1
-  python3 step_2_train_model.py \
-  --train_dataset "${train_dataset}" \
-  --valid_dataset "${valid_dataset}" \
-  --serialization_dir "${file_dir}" \
-  --config_file "${config_file}" \
-
-fi
-
-
-if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
-  $verbose && echo "stage 3: test model"
-  cd "${work_dir}" || exit 1
-  python3 step_3_evaluation.py \
-  --valid_dataset "${valid_dataset}" \
-  --model_dir "${file_dir}/best" \
-  --evaluation_audio_dir "${evaluation_audio_dir}" \
-  --limit "${limit}" \
-
-fi
-
-
-if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
-  $verbose && echo "stage 4: export model"
-  cd "${work_dir}" || exit 1
-  python3 step_5_export_models.py \
-  --vocabulary_dir "${vocabulary_dir}" \
-  --model_dir "${file_dir}/best" \
-  --serialization_dir "${file_dir}" \
-
-fi
-
-
-if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
-  $verbose && echo "stage 5: collect files"
-  cd "${work_dir}" || exit 1
-
-  mkdir -p ${final_model_dir}
-
-  cp "${file_dir}/best"/* "${final_model_dir}"
-  cp -r "${file_dir}/vocabulary" "${final_model_dir}"
-
-  cp "${file_dir}/evaluation.xlsx" "${final_model_dir}/evaluation.xlsx"
-
-  cp "${file_dir}/trace_model.zip" "${final_model_dir}/trace_model.zip"
-  cp "${file_dir}/trace_quant_model.zip" "${final_model_dir}/trace_quant_model.zip"
-  cp "${file_dir}/script_model.zip" "${final_model_dir}/script_model.zip"
-  cp "${file_dir}/script_quant_model.zip" "${final_model_dir}/script_quant_model.zip"
-
-  cd "${final_model_dir}/.." || exit 1;
-
-  if [ -e "${final_model_name}.zip" ]; then
-    rm -rf "${final_model_name}_backup.zip"
-    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
-  fi
-
-  zip -r "${final_model_name}.zip" "${final_model_name}"
-  rm -rf "${final_model_name}"
-
-fi
-
-
-if [ ${stage} -le 6 ] && [ ${stop_stage} -ge 6 ]; then
-  $verbose && echo "stage 6: clear file_dir"
-  cd "${work_dir}" || exit 1
-
-  rm -rf "${file_dir}";
-
-fi

examples/spectrum_dfnet_aishell/step_1_prepare_data.py

@@ -1,197 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import argparse
-import os
-from pathlib import Path
-import random
-import sys
-import shutil
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import pandas as pd
-from scipy.io import wavfile
-from tqdm import tqdm
-import librosa
-
-from project_settings import project_path
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--file_dir", default="./", type=str)
-
-    parser.add_argument(
-        "--noise_dir",
-        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
-        type=str
-    )
-    parser.add_argument(
-        "--speech_dir",
-        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
-        type=str
-    )
-
-    parser.add_argument("--train_dataset", default="train.xlsx", type=str)
-    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
-
-    parser.add_argument("--duration", default=2.0, type=float)
-    parser.add_argument("--min_snr_db", default=-10, type=float)
-    parser.add_argument("--max_snr_db", default=20, type=float)
-
-    parser.add_argument("--target_sample_rate", default=8000, type=int)
-
-    args = parser.parse_args()
-    return args
-
-
-def filename_generator(data_dir: str):
-    data_dir = Path(data_dir)
-    for filename in data_dir.glob("**/*.wav"):
-        yield filename.as_posix()
-
-
-def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000):
-    data_dir = Path(data_dir)
-    for filename in data_dir.glob("**/*.wav"):
-        signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
-        raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
-
-        if raw_duration < duration:
-            # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
-            continue
-        if signal.ndim != 1:
-            raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
-
-        signal_length = len(signal)
-        win_size = int(duration * sample_rate)
-        for begin in range(0, signal_length - win_size, win_size):
-            row = {
-                "filename": filename.as_posix(),
-                "raw_duration": round(raw_duration, 4),
-                "offset": round(begin / sample_rate, 4),
-                "duration": round(duration, 4),
-            }
-            yield row
-
-
-def get_dataset(args):
-    file_dir = Path(args.file_dir)
-    file_dir.mkdir(exist_ok=True)
-
-    noise_dir = Path(args.noise_dir)
-    speech_dir = Path(args.speech_dir)
-
-    noise_generator = target_second_signal_generator(
-        noise_dir.as_posix(),
-        duration=args.duration,
-        sample_rate=args.target_sample_rate
-    )
-    speech_generator = target_second_signal_generator(
-        speech_dir.as_posix(),
-        duration=args.duration,
-        sample_rate=args.target_sample_rate
-    )
-
-    dataset = list()
-
-    count = 0
-    process_bar = tqdm(desc="build dataset excel")
-    for noise, speech in zip(noise_generator, speech_generator):
-
-        noise_filename = noise["filename"]
-        noise_raw_duration = noise["raw_duration"]
-        noise_offset = noise["offset"]
-        noise_duration = noise["duration"]
-
-        speech_filename = speech["filename"]
-        speech_raw_duration = speech["raw_duration"]
-        speech_offset = speech["offset"]
-        speech_duration = speech["duration"]
-
-        random1 = random.random()
-        random2 = random.random()
-
-        row = {
-            "noise_filename": noise_filename,
-            "noise_raw_duration": noise_raw_duration,
-            "noise_offset": noise_offset,
-            "noise_duration": noise_duration,
-
-            "speech_filename": speech_filename,
-            "speech_raw_duration": speech_raw_duration,
-            "speech_offset": speech_offset,
-            "speech_duration": speech_duration,
-
-            "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
-
-            "random1": random1,
-            "random2": random2,
-            "flag": "TRAIN" if random2 < 0.8 else "TEST",
-        }
-        dataset.append(row)
-        count += 1
-        duration_seconds = count * args.duration
-        duration_hours = duration_seconds / 3600
-
-        process_bar.update(n=1)
-        process_bar.set_postfix({
-            # "duration_seconds": round(duration_seconds, 4),
-            "duration_hours": round(duration_hours, 4),
-
-        })
-
-    dataset = pd.DataFrame(dataset)
-    dataset = dataset.sort_values(by=["random1"], ascending=False)
-    dataset.to_excel(
-        file_dir / "dataset.xlsx",
-        index=False,
-    )
-    return
-
-
-
-def split_dataset(args):
-    """分割训练集, 测试集"""
-    file_dir = Path(args.file_dir)
-    file_dir.mkdir(exist_ok=True)
-
-    df = pd.read_excel(file_dir / "dataset.xlsx")
-
-    train = list()
-    test = list()
-
-    for i, row in df.iterrows():
-        flag = row["flag"]
-        if flag == "TRAIN":
-            train.append(row)
-        else:
-            test.append(row)
-
-    train = pd.DataFrame(train)
-    train.to_excel(
-        args.train_dataset,
-        index=False,
-        # encoding="utf_8_sig"
-    )
-    test = pd.DataFrame(test)
-    test.to_excel(
-        args.valid_dataset,
-        index=False,
-        # encoding="utf_8_sig"
-    )
-
-    return
-
-
-def main():
-    args = get_args()
-
-    get_dataset(args)
-    split_dataset(args)
-    return
-
-
-if __name__ == "__main__":
-    main()

examples/spectrum_dfnet_aishell/step_2_train_model.py

@@ -1,440 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-"""
-https://github.com/WenzheLiu-Speech/awesome-speech-enhancement
-"""
-import argparse
-import json
-import logging
-from logging.handlers import TimedRotatingFileHandler
-import os
-import platform
-from pathlib import Path
-import random
-import sys
-import shutil
-from typing import List
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-from torch.utils.data.dataloader import DataLoader
-import torchaudio
-from tqdm import tqdm
-
-from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
-from toolbox.torchaudio.models.spectrum_dfnet.configuration_spectrum_dfnet import SpectrumDfNetConfig
-from toolbox.torchaudio.models.spectrum_dfnet.modeling_spectrum_dfnet import SpectrumDfNetPretrainedModel
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--train_dataset", default="train.xlsx", type=str)
-    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
-
-    parser.add_argument("--max_epochs", default=100, type=int)
-
-    parser.add_argument("--batch_size", default=16, type=int)
-    parser.add_argument("--learning_rate", default=1e-4, type=float)
-    parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
-    parser.add_argument("--patience", default=5, type=int)
-    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
-    parser.add_argument("--seed", default=0, type=int)
-
-    parser.add_argument("--config_file", default="config.yaml", type=str)
-
-    args = parser.parse_args()
-    return args
-
-
-def logging_config(file_dir: str):
-    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
-
-    logging.basicConfig(format=fmt,
-                        datefmt="%m/%d/%Y %H:%M:%S",
-                        level=logging.INFO)
-    file_handler = TimedRotatingFileHandler(
-        filename=os.path.join(file_dir, "main.log"),
-        encoding="utf-8",
-        when="D",
-        interval=1,
-        backupCount=7
-    )
-    file_handler.setLevel(logging.INFO)
-    file_handler.setFormatter(logging.Formatter(fmt))
-    logger = logging.getLogger(__name__)
-    logger.addHandler(file_handler)
-
-    return logger
-
-
-class CollateFunction(object):
-    def __init__(self,
-                 n_fft: int = 512,
-                 win_length: int = 200,
-                 hop_length: int = 80,
-                 window_fn: str = "hamming",
-                 irm_beta: float = 1.0,
-                 epsilon: float = 1e-8,
-                 ):
-        self.n_fft = n_fft
-        self.win_length = win_length
-        self.hop_length = hop_length
-        self.window_fn = window_fn
-        self.irm_beta = irm_beta
-        self.epsilon = epsilon
-
-        self.complex_transform = torchaudio.transforms.Spectrogram(
-            n_fft=self.n_fft,
-            win_length=self.win_length,
-            hop_length=self.hop_length,
-            power=None,
-            window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
-        )
-        self.transform = torchaudio.transforms.Spectrogram(
-            n_fft=self.n_fft,
-            win_length=self.win_length,
-            hop_length=self.hop_length,
-            power=2.0,
-            window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
-        )
-
-    @staticmethod
-    def make_unfold_snr_db(x: torch.Tensor, n_time_steps: int = 3):
-        batch_size, channels, freq_dim, time_steps = x.shape
-
-        # kernel: [freq_dim, n_time_step]
-        kernel_size = (freq_dim, n_time_steps)
-
-        # pad
-        pad = n_time_steps // 2
-        x = torch.concat(tensors=[
-            x[:, :, :, :pad],
-            x,
-            x[:, :, :, -pad:],
-        ], dim=-1)
-
-        x = F.unfold(
-            input=x,
-            kernel_size=kernel_size,
-        )
-        # x shape: [batch_size, fold, time_steps]
-        return x
-
-    def __call__(self, batch: List[dict]):
-        speech_complex_spec_list = list()
-        mix_complex_spec_list = list()
-        speech_irm_list = list()
-        snr_db_list = list()
-        for sample in batch:
-            noise_wave: torch.Tensor = sample["noise_wave"]
-            speech_wave: torch.Tensor = sample["speech_wave"]
-            mix_wave: torch.Tensor = sample["mix_wave"]
-            # snr_db: float = sample["snr_db"]
-
-            noise_spec = self.transform.forward(noise_wave)
-            speech_spec = self.transform.forward(speech_wave)
-
-            speech_complex_spec = self.complex_transform.forward(speech_wave)
-            mix_complex_spec = self.complex_transform.forward(mix_wave)
-
-            # noise_irm = noise_spec / (noise_spec + speech_spec)
-            speech_irm = speech_spec / (noise_spec + speech_spec + self.epsilon)
-            speech_irm = torch.pow(speech_irm, self.irm_beta)
-
-            # noise_spec, speech_spec, mix_spec, speech_irm
-            # shape: [freq_dim, time_steps]
-
-            snr_db: torch.Tensor = 10 * torch.log10(
-                speech_spec / (noise_spec + self.epsilon)
-            )
-            snr_db = torch.clamp(snr_db, min=self.epsilon)
-
-            snr_db_ = torch.unsqueeze(snr_db, dim=0)
-            snr_db_ = torch.unsqueeze(snr_db_, dim=0)
-            snr_db_ = self.make_unfold_snr_db(snr_db_, n_time_steps=3)
-            snr_db_ = torch.squeeze(snr_db_, dim=0)
-            # snr_db_ shape: [fold, time_steps]
-
-            snr_db = torch.mean(snr_db_, dim=0, keepdim=True)
-            # snr_db shape: [1, time_steps]
-
-            speech_complex_spec_list.append(speech_complex_spec)
-            mix_complex_spec_list.append(mix_complex_spec)
-            speech_irm_list.append(speech_irm)
-            snr_db_list.append(snr_db)
-
-        speech_complex_spec_list = torch.stack(speech_complex_spec_list)
-        mix_complex_spec_list = torch.stack(mix_complex_spec_list)
-        speech_irm_list = torch.stack(speech_irm_list)
-        snr_db_list = torch.stack(snr_db_list)  # shape: (batch_size, time_steps, 1)
-
-        speech_complex_spec_list = speech_complex_spec_list[:, :-1, :]
-        mix_complex_spec_list = mix_complex_spec_list[:, :-1, :]
-        speech_irm_list = speech_irm_list[:, :-1, :]
-
-        # speech_complex_spec_list shape: [batch_size, freq_dim, time_steps]
-        # mix_complex_spec_list shape: [batch_size, freq_dim, time_steps]
-        # speech_irm_list shape: [batch_size, freq_dim, time_steps]
-        # snr_db shape: [batch_size, 1, time_steps]
-
-        # assert
-        if torch.any(torch.isnan(speech_complex_spec_list)) or torch.any(torch.isinf(speech_complex_spec_list)):
-            raise AssertionError("nan or inf in speech_complex_spec_list")
-        if torch.any(torch.isnan(mix_complex_spec_list)) or torch.any(torch.isinf(mix_complex_spec_list)):
-            raise AssertionError("nan or inf in mix_complex_spec_list")
-        if torch.any(torch.isnan(speech_irm_list)) or torch.any(torch.isinf(speech_irm_list)):
-            raise AssertionError("nan or inf in speech_irm_list")
-        if torch.any(torch.isnan(snr_db_list)) or torch.any(torch.isinf(snr_db_list)):
-            raise AssertionError("nan or inf in snr_db_list")
-
-        return speech_complex_spec_list, mix_complex_spec_list, speech_irm_list, snr_db_list
-
-
-collate_fn = CollateFunction()
-
-
-def main():
-    args = get_args()
-
-    serialization_dir = Path(args.serialization_dir)
-    serialization_dir.mkdir(parents=True, exist_ok=True)
-
-    logger = logging_config(serialization_dir)
-
-    random.seed(args.seed)
-    np.random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    logger.info("set seed: {}".format(args.seed))
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    n_gpu = torch.cuda.device_count()
-    logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
-
-    # datasets
-    logger.info("prepare datasets")
-    train_dataset = DenoiseExcelDataset(
-        excel_file=args.train_dataset,
-        expected_sample_rate=8000,
-        max_wave_value=32768.0,
-    )
-    valid_dataset = DenoiseExcelDataset(
-        excel_file=args.valid_dataset,
-        expected_sample_rate=8000,
-        max_wave_value=32768.0,
-    )
-    train_data_loader = DataLoader(
-        dataset=train_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
-        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
-        collate_fn=collate_fn,
-        pin_memory=False,
-        # prefetch_factor=64,
-    )
-    valid_data_loader = DataLoader(
-        dataset=valid_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
-        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
-        collate_fn=collate_fn,
-        pin_memory=False,
-        # prefetch_factor=64,
-    )
-
-    # models
-    logger.info(f"prepare models. config_file: {args.config_file}")
-    config = SpectrumDfNetConfig.from_pretrained(
-        pretrained_model_name_or_path=args.config_file,
-        # num_labels=vocabulary.get_vocab_size(namespace="labels")
-    )
-    model = SpectrumDfNetPretrainedModel(
-        config=config,
-    )
-    model.to(device)
-    model.train()
-
-    # optimizer
-    logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
-    param_optimizer = model.parameters()
-    optimizer = torch.optim.Adam(
-        param_optimizer,
-        lr=args.learning_rate,
-    )
-    # lr_scheduler = torch.optim.lr_scheduler.StepLR(
-    #     optimizer,
-    #     step_size=2000
-    # )
-    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-        optimizer,
-        milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
-    )
-
-    speech_mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-    irm_mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-    snr_mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-
-    # training loop
-    logger.info("training")
-
-    training_loss = 10000000000
-    evaluation_loss = 10000000000
-
-    model_list = list()
-    best_idx_epoch = None
-    best_metric = None
-    patience_count = 0
-
-    for idx_epoch in range(args.max_epochs):
-        total_loss = 0.
-        total_examples = 0.
-        progress_bar = tqdm(
-            total=len(train_data_loader),
-            desc="Training; epoch: {}".format(idx_epoch),
-        )
-
-        for batch in train_data_loader:
-            speech_complex_spec, mix_complex_spec, speech_irm, snr_db = batch
-            speech_complex_spec = speech_complex_spec.to(device)
-            mix_complex_spec = mix_complex_spec.to(device)
-            speech_irm_target = speech_irm.to(device)
-            snr_db_target = snr_db.to(device)
-
-            speech_spec_prediction, speech_irm_prediction, lsnr_prediction = model.forward(mix_complex_spec)
-            if torch.any(torch.isnan(speech_spec_prediction)) or torch.any(torch.isinf(speech_spec_prediction)):
-                raise AssertionError("nan or inf in speech_spec_prediction")
-            if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
-                raise AssertionError("nan or inf in speech_irm_prediction")
-            if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
-                raise AssertionError("nan or inf in lsnr_prediction")
-
-            speech_loss = speech_mse_loss.forward(speech_spec_prediction, torch.view_as_real(speech_complex_spec))
-            irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
-            snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
-
-            loss = speech_loss + irm_loss + snr_loss
-
-            total_loss += loss.item()
-            total_examples += mix_complex_spec.size(0)
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-            lr_scheduler.step()
-
-            training_loss = total_loss / total_examples
-            training_loss = round(training_loss, 4)
-
-            progress_bar.update(1)
-            progress_bar.set_postfix({
-                "training_loss": training_loss,
-            })
-
-        total_loss = 0.
-        total_examples = 0.
-        progress_bar = tqdm(
-            total=len(valid_data_loader),
-            desc="Evaluation; epoch: {}".format(idx_epoch),
-        )
-        for batch in valid_data_loader:
-            speech_complex_spec, mix_complex_spec, speech_irm, snr_db = batch
-            speech_complex_spec = speech_complex_spec.to(device)
-            mix_complex_spec = mix_complex_spec.to(device)
-            speech_irm_target = speech_irm.to(device)
-            snr_db_target = snr_db.to(device)
-
-            with torch.no_grad():
-                speech_spec_prediction, speech_irm_prediction, lsnr_prediction = model.forward(mix_complex_spec)
-                if torch.any(torch.isnan(speech_spec_prediction)) or torch.any(torch.isinf(speech_spec_prediction)):
-                    raise AssertionError("nan or inf in speech_spec_prediction")
-                if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
-                    raise AssertionError("nan or inf in speech_irm_prediction")
-                if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
-                    raise AssertionError("nan or inf in lsnr_prediction")
-
-                speech_loss = speech_mse_loss.forward(speech_spec_prediction, torch.view_as_real(speech_complex_spec))
-                irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
-                snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
-
-                loss = speech_loss + irm_loss + snr_loss
-
-                total_loss += loss.item()
-            total_examples += mix_complex_spec.size(0)
-
-            evaluation_loss = total_loss / total_examples
-            evaluation_loss = round(evaluation_loss, 4)
-
-            progress_bar.update(1)
-            progress_bar.set_postfix({
-                "evaluation_loss": evaluation_loss,
-            })
-
-        # save path
-        epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
-        epoch_dir.mkdir(parents=True, exist_ok=False)
-
-        # save models
-        model.save_pretrained(epoch_dir.as_posix())
-
-        model_list.append(epoch_dir)
-        if len(model_list) >= args.num_serialized_models_to_keep:
-            model_to_delete: Path = model_list.pop(0)
-            shutil.rmtree(model_to_delete.as_posix())
-
-        # save metric
-        if best_metric is None:
-            best_idx_epoch = idx_epoch
-            best_metric = evaluation_loss
-        elif evaluation_loss < best_metric:
-            best_idx_epoch = idx_epoch
-            best_metric = evaluation_loss
-        else:
-            pass
-
-        metrics = {
-            "idx_epoch": idx_epoch,
-            "best_idx_epoch": best_idx_epoch,
-            "training_loss": training_loss,
-            "evaluation_loss": evaluation_loss,
-            "learning_rate": optimizer.param_groups[0]["lr"],
-        }
-        metrics_filename = epoch_dir / "metrics_epoch.json"
-        with open(metrics_filename, "w", encoding="utf-8") as f:
-            json.dump(metrics, f, indent=4, ensure_ascii=False)
-
-        # save best
-        best_dir = serialization_dir / "best"
-        if best_idx_epoch == idx_epoch:
-            if best_dir.exists():
-                shutil.rmtree(best_dir)
-            shutil.copytree(epoch_dir, best_dir)
-
-        # early stop
-        early_stop_flag = False
-        if best_idx_epoch == idx_epoch:
-            patience_count = 0
-        else:
-            patience_count += 1
-        if patience_count >= args.patience:
-            early_stop_flag = True
-
-        # early stop
-        if early_stop_flag:
-            break
-    return
-
-
-if __name__ == '__main__':
-    main()

examples/spectrum_dfnet_aishell/step_3_evaluation.py

@@ -1,302 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import argparse
-import logging
-import os
-from pathlib import Path
-import sys
-import uuid
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import librosa
-import numpy as np
-import pandas as pd
-from scipy.io import wavfile
-import torch
-import torch.nn as nn
-import torchaudio
-from tqdm import tqdm
-
-from toolbox.torchaudio.models.spectrum_dfnet.modeling_spectrum_dfnet import SpectrumDfNetPretrainedModel
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
-    parser.add_argument("--model_dir", default="serialization_dir/best", type=str)
-    parser.add_argument("--evaluation_audio_dir", default="evaluation_audio_dir", type=str)
-
-    parser.add_argument("--limit", default=10, type=int)
-
-    args = parser.parse_args()
-    return args
-
-
-def logging_config():
-    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
-
-    logging.basicConfig(format=fmt,
-                        datefmt="%m/%d/%Y %H:%M:%S",
-                        level=logging.INFO)
-    stream_handler = logging.StreamHandler()
-    stream_handler.setLevel(logging.INFO)
-    stream_handler.setFormatter(logging.Formatter(fmt))
-
-    logger = logging.getLogger(__name__)
-
-    return logger
-
-
-def mix_speech_and_noise(speech: np.ndarray, noise: np.ndarray, snr_db: float):
-    l1 = len(speech)
-    l2 = len(noise)
-    l = min(l1, l2)
-    speech = speech[:l]
-    noise = noise[:l]
-
-    # np.float32, value between (-1, 1).
-
-    speech_power = np.mean(np.square(speech))
-    noise_power = speech_power / (10 ** (snr_db / 10))
-
-    noise_adjusted = np.sqrt(noise_power) * noise / np.sqrt(np.mean(noise ** 2))
-
-    noisy_signal = speech + noise_adjusted
-
-    return noisy_signal
-
-
-stft_power = torchaudio.transforms.Spectrogram(
-    n_fft=512,
-    win_length=200,
-    hop_length=80,
-    power=2.0,
-    window_fn=torch.hamming_window,
-)
-
-
-stft_complex = torchaudio.transforms.Spectrogram(
-    n_fft=512,
-    win_length=200,
-    hop_length=80,
-    power=None,
-    window_fn=torch.hamming_window,
-)
-
-
-istft = torchaudio.transforms.InverseSpectrogram(
-    n_fft=512,
-    win_length=200,
-    hop_length=80,
-    window_fn=torch.hamming_window,
-)
-
-
-def enhance(mix_spec_complex: torch.Tensor,
-            speech_spec_prediction: torch.Tensor,
-            speech_irm_prediction: torch.Tensor,
-            ):
-    mix_spec_complex = mix_spec_complex.detach().cpu()
-    speech_spec_prediction = speech_spec_prediction.detach().cpu()
-    speech_irm_prediction = speech_irm_prediction.detach().cpu()
-
-    mask_speech = speech_irm_prediction
-    mask_noise = 1.0 - speech_irm_prediction
-
-    speech_spec = mix_spec_complex * mask_speech
-    noise_spec = mix_spec_complex * mask_noise
-
-    # print(f"speech_spec_prediction: {speech_spec_prediction.shape}")
-    # print(f"noise_spec: {noise_spec.shape}")
-
-    speech_wave = istft.forward(speech_spec_prediction)
-    # speech_wave = istft.forward(speech_spec)
-    noise_wave = istft.forward(noise_spec)
-
-    return speech_wave, noise_wave
-
-
-def save_audios(noise_wave: torch.Tensor,
-                speech_wave: torch.Tensor,
-                mix_wave: torch.Tensor,
-                speech_wave_enhanced: torch.Tensor,
-                noise_wave_enhanced: torch.Tensor,
-                output_dir: str,
-                sample_rate: int = 8000,
-                ):
-    basename = uuid.uuid4().__str__()
-    output_dir = Path(output_dir) / basename
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    filename = output_dir / "noise_wave.wav"
-    torchaudio.save(filename, noise_wave, sample_rate)
-    filename = output_dir / "speech_wave.wav"
-    torchaudio.save(filename, speech_wave, sample_rate)
-    filename = output_dir / "mix_wave.wav"
-    torchaudio.save(filename, mix_wave, sample_rate)
-
-    filename = output_dir / "speech_wave_enhanced.wav"
-    torchaudio.save(filename, speech_wave_enhanced, sample_rate)
-    filename = output_dir / "noise_wave_enhanced.wav"
-    torchaudio.save(filename, noise_wave_enhanced, sample_rate)
-
-    return output_dir.as_posix()
-
-
-def main():
-    args = get_args()
-
-    logger = logging_config()
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    n_gpu = torch.cuda.device_count()
-    logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
-
-    logger.info("prepare model")
-    model = SpectrumDfNetPretrainedModel.from_pretrained(
-        pretrained_model_name_or_path=args.model_dir,
-    )
-    model.to(device)
-    model.eval()
-
-    # optimizer
-    logger.info("prepare loss_fn")
-    irm_mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-    snr_mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-
-    logger.info("read excel")
-    df = pd.read_excel(args.valid_dataset)
-
-    total_loss = 0.
-    total_examples = 0.
-    progress_bar = tqdm(total=len(df), desc="Evaluation")
-    for idx, row in df.iterrows():
-        noise_filename = row["noise_filename"]
-        noise_offset = row["noise_offset"]
-        noise_duration = row["noise_duration"]
-
-        speech_filename = row["speech_filename"]
-        speech_offset = row["speech_offset"]
-        speech_duration = row["speech_duration"]
-
-        snr_db = row["snr_db"]
-
-        noise_wave, _ = librosa.load(
-            noise_filename,
-            sr=8000,
-            offset=noise_offset,
-            duration=noise_duration,
-        )
-        speech_wave, _ = librosa.load(
-            speech_filename,
-            sr=8000,
-            offset=speech_offset,
-            duration=speech_duration,
-        )
-        mix_wave: np.ndarray = mix_speech_and_noise(
-            speech=speech_wave,
-            noise=noise_wave,
-            snr_db=snr_db,
-        )
-        noise_wave = torch.tensor(noise_wave, dtype=torch.float32)
-        speech_wave = torch.tensor(speech_wave, dtype=torch.float32)
-        mix_wave: torch.Tensor = torch.tensor(mix_wave, dtype=torch.float32)
-
-        noise_wave = noise_wave.unsqueeze(dim=0)
-        speech_wave = speech_wave.unsqueeze(dim=0)
-        mix_wave = mix_wave.unsqueeze(dim=0)
-
-        noise_spec: torch.Tensor = stft_power.forward(noise_wave)
-        speech_spec: torch.Tensor = stft_power.forward(speech_wave)
-        mix_spec: torch.Tensor = stft_power.forward(mix_wave)
-
-        speech_spec_complex: torch.Tensor = stft_complex.forward(speech_wave)
-        mix_spec_complex: torch.Tensor = stft_complex.forward(mix_wave)
-        # mix_spec_complex shape: [batch_size, freq_dim (257), time_steps, 2]
-
-        noise_spec = noise_spec[:, :-1, :]
-        speech_spec = speech_spec[:, :-1, :]
-        mix_spec = mix_spec[:, :-1, :]
-        speech_spec_complex = speech_spec_complex[:, :-1, :]
-        mix_spec_complex = mix_spec_complex[:, :-1, :]
-
-        speech_irm = speech_spec / (noise_spec + speech_spec)
-        speech_irm = torch.pow(speech_irm, 1.0)
-
-        snr_db: torch.Tensor = 10 * torch.log10(
-            speech_spec / (noise_spec + 1e-8)
-        )
-        snr_db = torch.clamp(snr_db, min=1e-8)
-        snr_db = torch.mean(snr_db, dim=1, keepdim=True)
-        # snr_db shape: [batch_size, 1, time_steps]
-
-        speech_spec_complex = speech_spec_complex.to(device)
-        mix_spec_complex = mix_spec_complex.to(device)
-        mix_spec = mix_spec.to(device)
-        speech_irm_target = speech_irm.to(device)
-        snr_db_target = snr_db.to(device)
-
-        with torch.no_grad():
-            speech_spec_prediction, speech_irm_prediction, lsnr_prediction = model.forward(mix_spec_complex)
-            speech_spec_prediction = torch.view_as_complex(speech_spec_prediction)
-
-            irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
-            # snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
-            # loss = irm_loss + 0.1 * snr_loss
-            loss = irm_loss
-
-        # mix_spec_complex shape: [batch_size, freq_dim (257), time_steps, 2]
-        # speech_irm_prediction shape: [batch_size, freq_dim (256), time_steps]
-        batch_size, _, time_steps = speech_irm_prediction.shape
-
-        mix_spec_complex = torch.concat(
-            [
-                mix_spec_complex,
-                torch.zeros(size=(batch_size, 1, time_steps), dtype=mix_spec_complex.dtype).to(device)
-            ],
-            dim=1,
-        )
-        speech_spec_prediction = torch.concat(
-            [
-                speech_spec_prediction,
-                torch.zeros(size=(batch_size, 1, time_steps), dtype=speech_spec_prediction.dtype).to(device)
-            ],
-            dim=1,
-        )
-        speech_irm_prediction = torch.concat(
-            [
-                speech_irm_prediction,
-                0.5 * torch.ones(size=(batch_size, 1, time_steps), dtype=speech_irm_prediction.dtype).to(device)
-            ],
-            dim=1,
-        )
-
-        # speech_irm_prediction shape: [batch_size, freq_dim (257), time_steps]
-        speech_wave_enhanced, noise_wave_enhanced = enhance(mix_spec_complex, speech_spec_prediction, speech_irm_prediction)
-        save_audios(noise_wave, speech_wave, mix_wave, speech_wave_enhanced, noise_wave_enhanced, args.evaluation_audio_dir)
-
-        total_loss += loss.item()
-        total_examples += mix_spec.size(0)
-
-        evaluation_loss = total_loss / total_examples
-        evaluation_loss = round(evaluation_loss, 4)
-
-        progress_bar.update(1)
-        progress_bar.set_postfix({
-            "evaluation_loss": evaluation_loss,
-        })
-
-        if idx > args.limit:
-            break
-
-    return
-
-
-if __name__ == '__main__':
-    main()

examples/spectrum_dfnet_aishell/yaml/config.yaml

@@ -1,53 +0,0 @@
-model_name: "spectrum_unet_irm"
-
-# spec
-sample_rate: 8000
-n_fft: 512
-win_length: 200
-hop_length: 80
-
-spec_bins: 256
-
-# model
-conv_channels: 64
-conv_kernel_size_input:
-  - 3
-  - 3
-conv_kernel_size_inner:
-  - 1
-  - 3
-conv_lookahead: 0
-
-convt_kernel_size_inner:
-  - 1
-  - 3
-
-embedding_hidden_size: 256
-encoder_combine_op: "concat"
-
-encoder_emb_skip_op: "none"
-encoder_emb_linear_groups: 16
-encoder_emb_hidden_size: 256
-
-encoder_linear_groups: 32
-
-lsnr_max: 30
-lsnr_min: -15
-norm_tau: 1.
-
-decoder_emb_num_layers: 3
-decoder_emb_skip_op: "none"
-decoder_emb_linear_groups: 16
-decoder_emb_hidden_size: 256
-
-df_decoder_hidden_size: 256
-df_num_layers: 2
-df_order: 5
-df_bins: 96
-df_gru_skip: "grouped_linear"
-df_decoder_linear_groups: 16
-df_pathway_kernel_size_t: 5
-df_lookahead: 2
-
-# runtime
-use_post_filter: true

examples/spectrum_unet_irm_aishell/run.sh

@@ -1,178 +0,0 @@
-#!/usr/bin/env bash
-
-: <<'END'
-
-
-sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name file_dir \
---noise_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "E:/programmer/asr_datasets/aishell/data_aishell/wav/train"
-
-
-sh run.sh --stage 1 --stop_stage 3 --system_version centos --file_folder_name file_dir \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
-
-sh run.sh --stage 3 --stop_stage 3 --system_version centos --file_folder_name file_dir \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
-
-
-END
-
-
-# params
-system_version="windows";
-verbose=true;
-stage=0 # start from 0 if you need to start from data preparation
-stop_stage=9
-
-work_dir="$(pwd)"
-file_folder_name=file_folder_name
-final_model_name=final_model_name
-config_file="yaml/config.yaml"
-limit=10
-
-noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
-speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
-
-nohup_name=nohup.out
-
-# model params
-batch_size=64
-max_epochs=200
-save_top_k=10
-patience=5
-
-
-# parse options
-while true; do
-  [ -z "${1:-}" ] && break;  # break if there are no arguments
-  case "$1" in
-    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
-      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
-      old_value="(eval echo \\$$name)";
-      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
-        was_bool=true;
-      else
-        was_bool=false;
-      fi
-
-      # Set the variable to the right value-- the escaped quotes make it work if
-      # the option had spaces, like --cmd "queue.pl -sync y"
-      eval "${name}=\"$2\"";
-
-      # Check that Boolean-valued arguments are really Boolean.
-      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
-        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
-        exit 1;
-      fi
-      shift 2;
-      ;;
-
-    *) break;
-  esac
-done
-
-file_dir="${work_dir}/${file_folder_name}"
-final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
-evaluation_audio_dir="${file_dir}/evaluation_audio"
-
-dataset="${file_dir}/dataset.xlsx"
-train_dataset="${file_dir}/train.xlsx"
-valid_dataset="${file_dir}/valid.xlsx"
-
-$verbose && echo "system_version: ${system_version}"
-$verbose && echo "file_folder_name: ${file_folder_name}"
-
-if [ $system_version == "windows" ]; then
-  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
-elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
-  #source /data/local/bin/nx_denoise/bin/activate
-  alias python3='/data/local/bin/nx_denoise/bin/python3'
-fi
-
-
-if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
-  $verbose && echo "stage 1: prepare data"
-  cd "${work_dir}" || exit 1
-  python3 step_1_prepare_data.py \
-  --file_dir "${file_dir}" \
-  --noise_dir "${noise_dir}" \
-  --speech_dir "${speech_dir}" \
-  --train_dataset "${train_dataset}" \
-  --valid_dataset "${valid_dataset}" \
-
-fi
-
-
-if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
-  $verbose && echo "stage 2: train model"
-  cd "${work_dir}" || exit 1
-  python3 step_2_train_model.py \
-  --train_dataset "${train_dataset}" \
-  --valid_dataset "${valid_dataset}" \
-  --serialization_dir "${file_dir}" \
-  --config_file "${config_file}" \
-
-fi
-
-
-if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
-  $verbose && echo "stage 3: test model"
-  cd "${work_dir}" || exit 1
-  python3 step_3_evaluation.py \
-  --valid_dataset "${valid_dataset}" \
-  --model_dir "${file_dir}/best" \
-  --evaluation_audio_dir "${evaluation_audio_dir}" \
-  --limit "${limit}" \
-
-fi
-
-
-if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
-  $verbose && echo "stage 4: export model"
-  cd "${work_dir}" || exit 1
-  python3 step_5_export_models.py \
-  --vocabulary_dir "${vocabulary_dir}" \
-  --model_dir "${file_dir}/best" \
-  --serialization_dir "${file_dir}" \
-
-fi
-
-
-if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
-  $verbose && echo "stage 5: collect files"
-  cd "${work_dir}" || exit 1
-
-  mkdir -p ${final_model_dir}
-
-  cp "${file_dir}/best"/* "${final_model_dir}"
-  cp -r "${file_dir}/vocabulary" "${final_model_dir}"
-
-  cp "${file_dir}/evaluation.xlsx" "${final_model_dir}/evaluation.xlsx"
-
-  cp "${file_dir}/trace_model.zip" "${final_model_dir}/trace_model.zip"
-  cp "${file_dir}/trace_quant_model.zip" "${final_model_dir}/trace_quant_model.zip"
-  cp "${file_dir}/script_model.zip" "${final_model_dir}/script_model.zip"
-  cp "${file_dir}/script_quant_model.zip" "${final_model_dir}/script_quant_model.zip"
-
-  cd "${final_model_dir}/.." || exit 1;
-
-  if [ -e "${final_model_name}.zip" ]; then
-    rm -rf "${final_model_name}_backup.zip"
-    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
-  fi
-
-  zip -r "${final_model_name}.zip" "${final_model_name}"
-  rm -rf "${final_model_name}"
-
-fi
-
-
-if [ ${stage} -le 6 ] && [ ${stop_stage} -ge 6 ]; then
-  $verbose && echo "stage 6: clear file_dir"
-  cd "${work_dir}" || exit 1
-
-  rm -rf "${file_dir}";
-
-fi

examples/spectrum_unet_irm_aishell/step_1_prepare_data.py

@@ -1,197 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import argparse
-import os
-from pathlib import Path
-import random
-import sys
-import shutil
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import pandas as pd
-from scipy.io import wavfile
-from tqdm import tqdm
-import librosa
-
-from project_settings import project_path
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--file_dir", default="./", type=str)
-
-    parser.add_argument(
-        "--noise_dir",
-        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
-        type=str
-    )
-    parser.add_argument(
-        "--speech_dir",
-        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
-        type=str
-    )
-
-    parser.add_argument("--train_dataset", default="train.xlsx", type=str)
-    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
-
-    parser.add_argument("--duration", default=2.0, type=float)
-    parser.add_argument("--min_snr_db", default=-10, type=float)
-    parser.add_argument("--max_snr_db", default=20, type=float)
-
-    parser.add_argument("--target_sample_rate", default=8000, type=int)
-
-    args = parser.parse_args()
-    return args
-
-
-def filename_generator(data_dir: str):
-    data_dir = Path(data_dir)
-    for filename in data_dir.glob("**/*.wav"):
-        yield filename.as_posix()
-
-
-def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000):
-    data_dir = Path(data_dir)
-    for filename in data_dir.glob("**/*.wav"):
-        signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
-        raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
-
-        if raw_duration < duration:
-            # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
-            continue
-        if signal.ndim != 1:
-            raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
-
-        signal_length = len(signal)
-        win_size = int(duration * sample_rate)
-        for begin in range(0, signal_length - win_size, win_size):
-            row = {
-                "filename": filename.as_posix(),
-                "raw_duration": round(raw_duration, 4),
-                "offset": round(begin / sample_rate, 4),
-                "duration": round(duration, 4),
-            }
-            yield row
-
-
-def get_dataset(args):
-    file_dir = Path(args.file_dir)
-    file_dir.mkdir(exist_ok=True)
-
-    noise_dir = Path(args.noise_dir)
-    speech_dir = Path(args.speech_dir)
-
-    noise_generator = target_second_signal_generator(
-        noise_dir.as_posix(),
-        duration=args.duration,
-        sample_rate=args.target_sample_rate
-    )
-    speech_generator = target_second_signal_generator(
-        speech_dir.as_posix(),
-        duration=args.duration,
-        sample_rate=args.target_sample_rate
-    )
-
-    dataset = list()
-
-    count = 0
-    process_bar = tqdm(desc="build dataset excel")
-    for noise, speech in zip(noise_generator, speech_generator):
-
-        noise_filename = noise["filename"]
-        noise_raw_duration = noise["raw_duration"]
-        noise_offset = noise["offset"]
-        noise_duration = noise["duration"]
-
-        speech_filename = speech["filename"]
-        speech_raw_duration = speech["raw_duration"]
-        speech_offset = speech["offset"]
-        speech_duration = speech["duration"]
-
-        random1 = random.random()
-        random2 = random.random()
-
-        row = {
-            "noise_filename": noise_filename,
-            "noise_raw_duration": noise_raw_duration,
-            "noise_offset": noise_offset,
-            "noise_duration": noise_duration,
-
-            "speech_filename": speech_filename,
-            "speech_raw_duration": speech_raw_duration,
-            "speech_offset": speech_offset,
-            "speech_duration": speech_duration,
-
-            "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
-
-            "random1": random1,
-            "random2": random2,
-            "flag": "TRAIN" if random2 < 0.8 else "TEST",
-        }
-        dataset.append(row)
-        count += 1
-        duration_seconds = count * args.duration
-        duration_hours = duration_seconds / 3600
-
-        process_bar.update(n=1)
-        process_bar.set_postfix({
-            # "duration_seconds": round(duration_seconds, 4),
-            "duration_hours": round(duration_hours, 4),
-
-        })
-
-    dataset = pd.DataFrame(dataset)
-    dataset = dataset.sort_values(by=["random1"], ascending=False)
-    dataset.to_excel(
-        file_dir / "dataset.xlsx",
-        index=False,
-    )
-    return
-
-
-
-def split_dataset(args):
-    """分割训练集, 测试集"""
-    file_dir = Path(args.file_dir)
-    file_dir.mkdir(exist_ok=True)
-
-    df = pd.read_excel(file_dir / "dataset.xlsx")
-
-    train = list()
-    test = list()
-
-    for i, row in df.iterrows():
-        flag = row["flag"]
-        if flag == "TRAIN":
-            train.append(row)
-        else:
-            test.append(row)
-
-    train = pd.DataFrame(train)
-    train.to_excel(
-        args.train_dataset,
-        index=False,
-        # encoding="utf_8_sig"
-    )
-    test = pd.DataFrame(test)
-    test.to_excel(
-        args.valid_dataset,
-        index=False,
-        # encoding="utf_8_sig"
-    )
-
-    return
-
-
-def main():
-    args = get_args()
-
-    get_dataset(args)
-    split_dataset(args)
-    return
-
-
-if __name__ == "__main__":
-    main()

examples/spectrum_unet_irm_aishell/step_2_train_model.py

@@ -1,420 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-"""
-https://github.com/WenzheLiu-Speech/awesome-speech-enhancement
-"""
-import argparse
-import json
-import logging
-from logging.handlers import TimedRotatingFileHandler
-import os
-import platform
-from pathlib import Path
-import random
-import sys
-import shutil
-from typing import List
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-from torch.utils.data.dataloader import DataLoader
-import torchaudio
-from tqdm import tqdm
-
-from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
-from toolbox.torchaudio.models.spectrum_unet_irm.configuration_specturm_unet_irm import SpectrumUnetIRMConfig
-from toolbox.torchaudio.models.spectrum_unet_irm.modeling_spectrum_unet_irm import SpectrumUnetIRMPretrainedModel
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--train_dataset", default="train.xlsx", type=str)
-    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
-
-    parser.add_argument("--max_epochs", default=100, type=int)
-
-    parser.add_argument("--batch_size", default=64, type=int)
-    parser.add_argument("--learning_rate", default=1e-4, type=float)
-    parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
-    parser.add_argument("--patience", default=5, type=int)
-    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
-    parser.add_argument("--seed", default=0, type=int)
-
-    parser.add_argument("--config_file", default="config.yaml", type=str)
-
-    args = parser.parse_args()
-    return args
-
-
-def logging_config(file_dir: str):
-    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
-
-    logging.basicConfig(format=fmt,
-                        datefmt="%m/%d/%Y %H:%M:%S",
-                        level=logging.INFO)
-    file_handler = TimedRotatingFileHandler(
-        filename=os.path.join(file_dir, "main.log"),
-        encoding="utf-8",
-        when="D",
-        interval=1,
-        backupCount=7
-    )
-    file_handler.setLevel(logging.INFO)
-    file_handler.setFormatter(logging.Formatter(fmt))
-    logger = logging.getLogger(__name__)
-    logger.addHandler(file_handler)
-
-    return logger
-
-
-class CollateFunction(object):
-    def __init__(self,
-                 n_fft: int = 512,
-                 win_length: int = 200,
-                 hop_length: int = 80,
-                 window_fn: str = "hamming",
-                 irm_beta: float = 1.0,
-                 epsilon: float = 1e-8,
-                 ):
-        self.n_fft = n_fft
-        self.win_length = win_length
-        self.hop_length = hop_length
-        self.window_fn = window_fn
-        self.irm_beta = irm_beta
-        self.epsilon = epsilon
-
-        self.transform = torchaudio.transforms.Spectrogram(
-            n_fft=self.n_fft,
-            win_length=self.win_length,
-            hop_length=self.hop_length,
-            power=2.0,
-            window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
-        )
-
-    @staticmethod
-    def make_unfold_snr_db(x: torch.Tensor, n_time_steps: int = 3):
-        batch_size, channels, freq_dim, time_steps = x.shape
-
-        # kernel: [freq_dim, n_time_step]
-        kernel_size = (freq_dim, n_time_steps)
-
-        # pad
-        pad = n_time_steps // 2
-        x = torch.concat(tensors=[
-            x[:, :, :, :pad],
-            x,
-            x[:, :, :, -pad:],
-        ], dim=-1)
-
-        x = F.unfold(
-            input=x,
-            kernel_size=kernel_size,
-        )
-        # x shape: [batch_size, fold, time_steps]
-        return x
-
-    def __call__(self, batch: List[dict]):
-        mix_spec_list = list()
-        speech_irm_list = list()
-        snr_db_list = list()
-        for sample in batch:
-            noise_wave: torch.Tensor = sample["noise_wave"]
-            speech_wave: torch.Tensor = sample["speech_wave"]
-            mix_wave: torch.Tensor = sample["mix_wave"]
-            # snr_db: float = sample["snr_db"]
-
-            noise_spec = self.transform.forward(noise_wave)
-            speech_spec = self.transform.forward(speech_wave)
-            mix_spec = self.transform.forward(mix_wave)
-
-            # noise_irm = noise_spec / (noise_spec + speech_spec)
-            speech_irm = speech_spec / (noise_spec + speech_spec + self.epsilon)
-            speech_irm = torch.pow(speech_irm, self.irm_beta)
-
-            # noise_spec, speech_spec, mix_spec, speech_irm
-            # shape: [freq_dim, time_steps]
-
-            snr_db: torch.Tensor = 10 * torch.log10(
-                speech_spec / (noise_spec + self.epsilon)
-            )
-            snr_db = torch.clamp(snr_db, min=self.epsilon)
-
-            snr_db_ = torch.unsqueeze(snr_db, dim=0)
-            snr_db_ = torch.unsqueeze(snr_db_, dim=0)
-            snr_db_ = self.make_unfold_snr_db(snr_db_, n_time_steps=3)
-            snr_db_ = torch.squeeze(snr_db_, dim=0)
-            # snr_db_ shape: [fold, time_steps]
-
-            snr_db = torch.mean(snr_db_, dim=0, keepdim=True)
-            # snr_db shape: [1, time_steps]
-
-            mix_spec_list.append(mix_spec)
-            speech_irm_list.append(speech_irm)
-            snr_db_list.append(snr_db)
-
-        mix_spec_list = torch.stack(mix_spec_list)
-        speech_irm_list = torch.stack(speech_irm_list)
-        snr_db_list = torch.stack(snr_db_list)  # shape: (batch_size, time_steps, 1)
-
-        mix_spec_list = mix_spec_list[:, :-1, :]
-        speech_irm_list = speech_irm_list[:, :-1, :]
-
-        # mix_spec_list shape: [batch_size, freq_dim, time_steps]
-        # speech_irm_list shape: [batch_size, freq_dim, time_steps]
-        # snr_db shape: [batch_size, 1, time_steps]
-
-        # assert
-        if torch.any(torch.isnan(mix_spec_list)) or torch.any(torch.isinf(mix_spec_list)):
-            raise AssertionError("nan or inf in mix_spec_list")
-        if torch.any(torch.isnan(speech_irm_list)) or torch.any(torch.isinf(speech_irm_list)):
-            raise AssertionError("nan or inf in speech_irm_list")
-        if torch.any(torch.isnan(snr_db_list)) or torch.any(torch.isinf(snr_db_list)):
-            raise AssertionError("nan or inf in snr_db_list")
-
-        return mix_spec_list, speech_irm_list, snr_db_list
-
-
-collate_fn = CollateFunction()
-
-
-def main():
-    args = get_args()
-
-    serialization_dir = Path(args.serialization_dir)
-    serialization_dir.mkdir(parents=True, exist_ok=True)
-
-    logger = logging_config(serialization_dir)
-
-    random.seed(args.seed)
-    np.random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    logger.info("set seed: {}".format(args.seed))
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    n_gpu = torch.cuda.device_count()
-    logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
-
-    # datasets
-    logger.info("prepare datasets")
-    train_dataset = DenoiseExcelDataset(
-        excel_file=args.train_dataset,
-        expected_sample_rate=8000,
-        max_wave_value=32768.0,
-    )
-    valid_dataset = DenoiseExcelDataset(
-        excel_file=args.valid_dataset,
-        expected_sample_rate=8000,
-        max_wave_value=32768.0,
-    )
-    train_data_loader = DataLoader(
-        dataset=train_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
-        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
-        collate_fn=collate_fn,
-        pin_memory=False,
-        # prefetch_factor=64,
-    )
-    valid_data_loader = DataLoader(
-        dataset=valid_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
-        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
-        collate_fn=collate_fn,
-        pin_memory=False,
-        # prefetch_factor=64,
-    )
-
-    # models
-    logger.info(f"prepare models. config_file: {args.config_file}")
-    config = SpectrumUnetIRMConfig.from_pretrained(
-        pretrained_model_name_or_path=args.config_file,
-        # num_labels=vocabulary.get_vocab_size(namespace="labels")
-    )
-    model = SpectrumUnetIRMPretrainedModel(
-        config=config,
-    )
-    model.to(device)
-    model.train()
-
-    # optimizer
-    logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
-    param_optimizer = model.parameters()
-    optimizer = torch.optim.Adam(
-        param_optimizer,
-        lr=args.learning_rate,
-    )
-    # lr_scheduler = torch.optim.lr_scheduler.StepLR(
-    #     optimizer,
-    #     step_size=2000
-    # )
-    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-        optimizer,
-        milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
-    )
-    irm_mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-    snr_mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-
-    # training loop
-    logger.info("training")
-
-    training_loss = 10000000000
-    evaluation_loss = 10000000000
-
-    model_list = list()
-    best_idx_epoch = None
-    best_metric = None
-    patience_count = 0
-
-    for idx_epoch in range(args.max_epochs):
-        total_loss = 0.
-        total_examples = 0.
-        progress_bar = tqdm(
-            total=len(train_data_loader),
-            desc="Training; epoch: {}".format(idx_epoch),
-        )
-
-        for batch in train_data_loader:
-            mix_spec, speech_irm, snr_db = batch
-            mix_spec = mix_spec.to(device)
-            speech_irm_target = speech_irm.to(device)
-            snr_db_target = snr_db.to(device)
-
-            speech_irm_prediction, lsnr_prediction = model.forward(mix_spec)
-            if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
-                raise AssertionError("nan or inf in speech_irm_prediction")
-            if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
-                raise AssertionError("nan or inf in lsnr_prediction")
-            irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
-            lsnr_prediction = (lsnr_prediction - config.lsnr_min) / (config.lsnr_max - config.lsnr_min)
-            if torch.max(lsnr_prediction) > 1 or torch.min(lsnr_prediction) < 0:
-                raise AssertionError(f"expected lsnr_prediction between 0 and 1.")
-            snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
-            if torch.any(torch.isnan(snr_loss)) or torch.any(torch.isinf(snr_loss)):
-                raise AssertionError("nan or inf in snr_loss")
-            # loss = irm_loss + 0.1 * snr_loss
-            loss = 10.0 * irm_loss + 0.05 * snr_loss
-            # loss = irm_loss
-
-            total_loss += loss.item()
-            total_examples += mix_spec.size(0)
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-            lr_scheduler.step()
-
-            training_loss = total_loss / total_examples
-            training_loss = round(training_loss, 4)
-
-            progress_bar.update(1)
-            progress_bar.set_postfix({
-                "training_loss": training_loss,
-            })
-
-        total_loss = 0.
-        total_examples = 0.
-        progress_bar = tqdm(
-            total=len(valid_data_loader),
-            desc="Evaluation; epoch: {}".format(idx_epoch),
-        )
-        for batch in valid_data_loader:
-            mix_spec, speech_irm, snr_db = batch
-            mix_spec = mix_spec.to(device)
-            speech_irm_target = speech_irm.to(device)
-            snr_db_target = snr_db.to(device)
-
-            with torch.no_grad():
-                speech_irm_prediction, lsnr_prediction = model.forward(mix_spec)
-                if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
-                    raise AssertionError("nan or inf in speech_irm_prediction")
-                if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
-                    raise AssertionError("nan or inf in lsnr_prediction")
-                irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
-                lsnr_prediction = (lsnr_prediction - config.lsnr_min) / (config.lsnr_max - config.lsnr_min)
-                if torch.max(lsnr_prediction) > 1 or torch.min(lsnr_prediction) < 0:
-                    raise AssertionError(f"expected lsnr_prediction between 0 and 1.")
-                snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
-                # loss = irm_loss + 0.1 * snr_loss
-                loss = 10.0 * irm_loss + 0.05 * snr_loss
-                # loss = irm_loss
-
-                total_loss += loss.item()
-            total_examples += mix_spec.size(0)
-
-            evaluation_loss = total_loss / total_examples
-            evaluation_loss = round(evaluation_loss, 4)
-
-            progress_bar.update(1)
-            progress_bar.set_postfix({
-                "evaluation_loss": evaluation_loss,
-            })
-
-        # save path
-        epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
-        epoch_dir.mkdir(parents=True, exist_ok=False)
-
-        # save models
-        model.save_pretrained(epoch_dir.as_posix())
-
-        model_list.append(epoch_dir)
-        if len(model_list) >= args.num_serialized_models_to_keep:
-            model_to_delete: Path = model_list.pop(0)
-            shutil.rmtree(model_to_delete.as_posix())
-
-        # save metric
-        if best_metric is None:
-            best_idx_epoch = idx_epoch
-            best_metric = evaluation_loss
-        elif evaluation_loss < best_metric:
-            best_idx_epoch = idx_epoch
-            best_metric = evaluation_loss
-        else:
-            pass
-
-        metrics = {
-            "idx_epoch": idx_epoch,
-            "best_idx_epoch": best_idx_epoch,
-            "training_loss": training_loss,
-            "evaluation_loss": evaluation_loss,
-            "learning_rate": optimizer.param_groups[0]["lr"],
-        }
-        metrics_filename = epoch_dir / "metrics_epoch.json"
-        with open(metrics_filename, "w", encoding="utf-8") as f:
-            json.dump(metrics, f, indent=4, ensure_ascii=False)
-
-        # save best
-        best_dir = serialization_dir / "best"
-        if best_idx_epoch == idx_epoch:
-            if best_dir.exists():
-                shutil.rmtree(best_dir)
-            shutil.copytree(epoch_dir, best_dir)
-
-        # early stop
-        early_stop_flag = False
-        if best_idx_epoch == idx_epoch:
-            patience_count = 0
-        else:
-            patience_count += 1
-        if patience_count >= args.patience:
-            early_stop_flag = True
-
-        # early stop
-        if early_stop_flag:
-            break
-    return
-
-
-if __name__ == '__main__':
-    main()

examples/spectrum_unet_irm_aishell/step_3_evaluation.py

@@ -1,270 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import argparse
-import logging
-import os
-from pathlib import Path
-import sys
-import uuid
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import librosa
-import numpy as np
-import pandas as pd
-from scipy.io import wavfile
-import torch
-import torch.nn as nn
-import torchaudio
-from tqdm import tqdm
-
-from toolbox.torchaudio.models.spectrum_unet_irm.modeling_spectrum_unet_irm import SpectrumUnetIRMPretrainedModel
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
-    parser.add_argument("--model_dir", default="serialization_dir/best", type=str)
-    parser.add_argument("--evaluation_audio_dir", default="evaluation_audio_dir", type=str)
-
-    parser.add_argument("--limit", default=10, type=int)
-
-    args = parser.parse_args()
-    return args
-
-
-def logging_config():
-    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
-
-    logging.basicConfig(format=fmt,
-                        datefmt="%m/%d/%Y %H:%M:%S",
-                        level=logging.INFO)
-    stream_handler = logging.StreamHandler()
-    stream_handler.setLevel(logging.INFO)
-    stream_handler.setFormatter(logging.Formatter(fmt))
-
-    logger = logging.getLogger(__name__)
-
-    return logger
-
-
-def mix_speech_and_noise(speech: np.ndarray, noise: np.ndarray, snr_db: float):
-    l1 = len(speech)
-    l2 = len(noise)
-    l = min(l1, l2)
-    speech = speech[:l]
-    noise = noise[:l]
-
-    # np.float32, value between (-1, 1).
-
-    speech_power = np.mean(np.square(speech))
-    noise_power = speech_power / (10 ** (snr_db / 10))
-
-    noise_adjusted = np.sqrt(noise_power) * noise / np.sqrt(np.mean(noise ** 2))
-
-    noisy_signal = speech + noise_adjusted
-
-    return noisy_signal
-
-
-stft_power = torchaudio.transforms.Spectrogram(
-    n_fft=512,
-    win_length=200,
-    hop_length=80,
-    power=2.0,
-    window_fn=torch.hamming_window,
-)
-
-
-stft_complex = torchaudio.transforms.Spectrogram(
-    n_fft=512,
-    win_length=200,
-    hop_length=80,
-    power=None,
-    window_fn=torch.hamming_window,
-)
-
-
-istft = torchaudio.transforms.InverseSpectrogram(
-    n_fft=512,
-    win_length=200,
-    hop_length=80,
-    window_fn=torch.hamming_window,
-)
-
-
-def enhance(mix_spec_complex: torch.Tensor, speech_irm_prediction: torch.Tensor):
-    mix_spec_complex = mix_spec_complex.detach().cpu()
-    speech_irm_prediction = speech_irm_prediction.detach().cpu()
-
-    mask_speech = speech_irm_prediction
-    mask_noise = 1.0 - speech_irm_prediction
-
-    speech_spec = mix_spec_complex * mask_speech
-    noise_spec = mix_spec_complex * mask_noise
-
-    speech_wave = istft.forward(speech_spec)
-    noise_wave = istft.forward(noise_spec)
-
-    return speech_wave, noise_wave
-
-
-def save_audios(noise_wave: torch.Tensor,
-                speech_wave: torch.Tensor,
-                mix_wave: torch.Tensor,
-                speech_wave_enhanced: torch.Tensor,
-                noise_wave_enhanced: torch.Tensor,
-                output_dir: str,
-                sample_rate: int = 8000,
-                ):
-    basename = uuid.uuid4().__str__()
-    output_dir = Path(output_dir) / basename
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    filename = output_dir / "noise_wave.wav"
-    torchaudio.save(filename, noise_wave, sample_rate)
-    filename = output_dir / "speech_wave.wav"
-    torchaudio.save(filename, speech_wave, sample_rate)
-    filename = output_dir / "mix_wave.wav"
-    torchaudio.save(filename, mix_wave, sample_rate)
-
-    filename = output_dir / "speech_wave_enhanced.wav"
-    torchaudio.save(filename, speech_wave_enhanced, sample_rate)
-    filename = output_dir / "noise_wave_enhanced.wav"
-    torchaudio.save(filename, noise_wave_enhanced, sample_rate)
-
-    return output_dir.as_posix()
-
-
-def main():
-    args = get_args()
-
-    logger = logging_config()
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    n_gpu = torch.cuda.device_count()
-    logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
-
-    logger.info("prepare model")
-    model = SpectrumUnetIRMPretrainedModel.from_pretrained(
-        pretrained_model_name_or_path=args.model_dir,
-    )
-    model.to(device)
-    model.eval()
-
-    # optimizer
-    logger.info("prepare loss_fn")
-    irm_mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-    snr_mse_loss = nn.MSELoss(
-        reduction="mean",
-    )
-
-    logger.info("read excel")
-    df = pd.read_excel(args.valid_dataset)
-
-    total_loss = 0.
-    total_examples = 0.
-    progress_bar = tqdm(total=len(df), desc="Evaluation")
-    for idx, row in df.iterrows():
-        noise_filename = row["noise_filename"]
-        noise_offset = row["noise_offset"]
-        noise_duration = row["noise_duration"]
-
-        speech_filename = row["speech_filename"]
-        speech_offset = row["speech_offset"]
-        speech_duration = row["speech_duration"]
-
-        snr_db = row["snr_db"]
-
-        noise_wave, _ = librosa.load(
-            noise_filename,
-            sr=8000,
-            offset=noise_offset,
-            duration=noise_duration,
-        )
-        speech_wave, _ = librosa.load(
-            speech_filename,
-            sr=8000,
-            offset=speech_offset,
-            duration=speech_duration,
-        )
-        mix_wave: np.ndarray = mix_speech_and_noise(
-            speech=speech_wave,
-            noise=noise_wave,
-            snr_db=snr_db,
-        )
-        noise_wave = torch.tensor(noise_wave, dtype=torch.float32)
-        speech_wave = torch.tensor(speech_wave, dtype=torch.float32)
-        mix_wave: torch.Tensor = torch.tensor(mix_wave, dtype=torch.float32)
-
-        noise_wave = noise_wave.unsqueeze(dim=0)
-        speech_wave = speech_wave.unsqueeze(dim=0)
-        mix_wave = mix_wave.unsqueeze(dim=0)
-
-        noise_spec: torch.Tensor = stft_power.forward(noise_wave)
-        speech_spec: torch.Tensor = stft_power.forward(speech_wave)
-        mix_spec: torch.Tensor = stft_power.forward(mix_wave)
-
-        noise_spec = noise_spec[:, :-1, :]
-        speech_spec = speech_spec[:, :-1, :]
-        mix_spec = mix_spec[:, :-1, :]
-
-        mix_spec_complex: torch.Tensor = stft_complex.forward(mix_wave)
-        # mix_spec_complex shape: [batch_size, freq_dim (257), time_steps, 2]
-
-        speech_irm = speech_spec / (noise_spec + speech_spec)
-        speech_irm = torch.pow(speech_irm, 1.0)
-
-        snr_db: torch.Tensor = 10 * torch.log10(
-            speech_spec / (noise_spec + 1e-8)
-        )
-        snr_db = torch.mean(snr_db, dim=1, keepdim=True)
-        # snr_db shape: [batch_size, 1, time_steps]
-
-        mix_spec = mix_spec.to(device)
-        speech_irm_target = speech_irm.to(device)
-        snr_db_target = snr_db.to(device)
-
-        with torch.no_grad():
-            speech_irm_prediction, lsnr_prediction = model.forward(mix_spec)
-            irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
-            # snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
-            # loss = irm_loss + 0.1 * snr_loss
-            loss = irm_loss
-
-        # mix_spec_complex shape: [batch_size, freq_dim (257), time_steps, 2]
-        # speech_irm_prediction shape: [batch_size, freq_dim (256), time_steps]
-        batch_size, _, time_steps = speech_irm_prediction.shape
-        speech_irm_prediction = torch.concat(
-            [
-                speech_irm_prediction,
-                0.5*torch.ones(size=(batch_size, 1, time_steps), dtype=speech_irm_prediction.dtype).to(device)
-            ],
-            dim=1,
-        )
-        # speech_irm_prediction shape: [batch_size, freq_dim (257), time_steps]
-        speech_wave_enhanced, noise_wave_enhanced = enhance(mix_spec_complex, speech_irm_prediction)
-        save_audios(noise_wave, speech_wave, mix_wave, speech_wave_enhanced, noise_wave_enhanced, args.evaluation_audio_dir)
-
-        total_loss += loss.item()
-        total_examples += mix_spec.size(0)
-
-        evaluation_loss = total_loss / total_examples
-        evaluation_loss = round(evaluation_loss, 4)
-
-        progress_bar.update(1)
-        progress_bar.set_postfix({
-            "evaluation_loss": evaluation_loss,
-        })
-
-        if idx > args.limit:
-            break
-
-    return
-
-
-if __name__ == '__main__':
-    main()

examples/spectrum_unet_irm_aishell/yaml/config.yaml

@@ -1,38 +0,0 @@
-model_name: "spectrum_unet_irm"
-
-# spec
-sample_rate: 8000
-n_fft: 512
-win_length: 200
-hop_length: 80
-
-spec_bins: 256
-
-# model
-conv_channels: 64
-conv_kernel_size_input:
-  - 3
-  - 3
-conv_kernel_size_inner:
-  - 1
-  - 3
-conv_lookahead: 0
-
-convt_kernel_size_inner:
-  - 1
-  - 3
-
-encoder_emb_skip_op: "none"
-encoder_emb_linear_groups: 16
-encoder_emb_hidden_size: 256
-
-lsnr_max: 30
-lsnr_min: -15
-
-decoder_emb_num_layers: 3
-decoder_emb_skip_op: "none"
-decoder_emb_linear_groups: 16
-decoder_emb_hidden_size: 256
-
-# runtime
-use_post_filter: true

main.py

@@ -1,7 +1,7 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 """
-docker build -t denoise:v20250609_1919 .
+docker build -t denoise:v20250626_1616 .
 docker stop denoise_7865 && docker rm denoise_7865
 docker run -itd \
 --name denoise_7865 \

toolbox/torch/utils/data/dataset/mp3_to_wav_jsonl_dataset.py

@@ -1,197 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import json
-import os
-import random
-from typing import List
-from pathlib import Path
-import tempfile
-import uuid
-
-from pydub import AudioSegment
-from scipy.io import wavfile
-import librosa
-import numpy as np
-import torch
-from torch.utils.data import Dataset, IterableDataset
-
-
-class Mp3ToWavJsonlDataset(IterableDataset):
-    def __init__(self,
-                 jsonl_file: str,
-                 expected_sample_rate: int,
-                 resample: bool = False,
-                 max_wave_value: float = 1.0,
-                 buffer_size: int = 1000,
-                 eps: float = 1e-8,
-                 skip: int = 0,
-                 ):
-        self.jsonl_file = jsonl_file
-        self.expected_sample_rate = expected_sample_rate
-        self.resample = resample
-        self.max_wave_value = max_wave_value
-        self.eps = eps
-        self.skip = skip
-
-        self.buffer_size = buffer_size
-        self.buffer_samples: List[dict] = list()
-
-    def __iter__(self):
-        self.buffer_samples = list()
-
-        iterable_source = self.iterable_source()
-
-        try:
-            for _ in range(self.skip):
-                next(iterable_source)
-        except StopIteration:
-            pass
-
-        # 初始填充缓冲区
-        try:
-            for _ in range(self.buffer_size):
-                self.buffer_samples.append(next(iterable_source))
-        except StopIteration:
-            pass
-
-        # 动态替换逻辑
-        while True:
-            try:
-                item = next(iterable_source)
-                # 随机替换缓冲区元素
-                replace_idx = random.randint(0, len(self.buffer_samples) - 1)
-                sample = self.buffer_samples[replace_idx]
-                self.buffer_samples[replace_idx] = item
-                yield self.convert_sample(sample)
-            except StopIteration:
-                break
-
-        # 清空剩余元素
-        random.shuffle(self.buffer_samples)
-        for sample in self.buffer_samples:
-            yield self.convert_sample(sample)
-
-    def iterable_source(self):
-        last_sample = None
-        with open(self.jsonl_file, "r", encoding="utf-8") as f:
-            for row in f:
-                row = json.loads(row)
-                filename = row["filename"]
-                raw_duration = row["raw_duration"]
-                offset = row["offset"]
-                duration = row["duration"]
-
-                sample = {
-                    "filename": filename,
-                    "raw_duration": raw_duration,
-                    "offset": offset,
-                    "duration": duration,
-                }
-                if last_sample is None:
-                    last_sample = sample
-                    continue
-                yield sample
-        yield last_sample
-
-    def convert_sample(self, sample: dict):
-        filename = sample["filename"]
-        offset = sample["offset"]
-        duration = sample["duration"]
-
-        wav_waveform = self.filename_to_waveform(filename, offset, duration)
-        mp3_waveform = self.filename_to_mp3_waveform(filename, offset, duration)
-
-        if wav_waveform.shape != mp3_waveform.shape:
-            raise AssertionError(f"wav_waveform: {wav_waveform.shape}, mp3_waveform: {mp3_waveform.shape}")
-
-        result = {
-            "mp3_waveform": mp3_waveform,
-            "wav_waveform": wav_waveform,
-        }
-        return result
-
-    @staticmethod
-    def filename_to_waveform(filename: str, offset: float, duration: float, expected_sample_rate: int = 8000):
-        try:
-            waveform, sample_rate = librosa.load(
-                filename,
-                sr=expected_sample_rate,
-                offset=offset,
-                duration=duration,
-            )
-        except ValueError as e:
-            print(f"load failed. error type: {type(e)}, error text: {str(e)}, filename: {filename}")
-            raise e
-        waveform = torch.tensor(waveform, dtype=torch.float32)
-        return waveform
-
-    @staticmethod
-    def get_temporary_file(suffix: str = ".wav"):
-        temp_audio_dir = Path(tempfile.gettempdir()) / "mp3_to_wav_jsonl_dataset"
-        temp_audio_dir.mkdir(parents=True, exist_ok=True)
-        filename = temp_audio_dir / f"{uuid.uuid4()}{suffix}"
-        filename = filename.as_posix()
-        return filename
-
-    @staticmethod
-    def filename_to_mp3_waveform(filename: str, offset: float, duration: float, expected_sample_rate: int = 8000):
-        try:
-            waveform, sample_rate = librosa.load(
-                filename,
-                sr=expected_sample_rate,
-                offset=offset,
-                duration=duration,
-            )
-            waveform = np.array(waveform * (1 << 15), dtype=np.int16)
-        except ValueError as e:
-            print(f"load failed. error type: {type(e)}, error text: {str(e)}, filename: {filename}")
-            raise e
-
-        wav_temporary_file = Mp3ToWavJsonlDataset.get_temporary_file(suffix=".wav")
-        wavfile.write(
-            wav_temporary_file,
-            rate=sample_rate,
-            data=waveform,
-        )
-
-        mp3_temporary_file = Mp3ToWavJsonlDataset.get_temporary_file(suffix=".mp3")
-
-        audio = AudioSegment.from_wav(wav_temporary_file)
-        audio.export(mp3_temporary_file,
-                     format="mp3",
-                     bitrate="64k",  # 8kHz建议使用64kbps
-                     # parameters=["-ar", "8000"]
-                     parameters=["-ar", f"{expected_sample_rate}"]
-                     )
-
-        try:
-            waveform, sample_rate = librosa.load(mp3_temporary_file, sr=expected_sample_rate)
-        except ValueError as e:
-            print(f"load failed. error type: {type(e)}, error text: {str(e)}, filename: {filename}")
-            raise e
-
-        os.remove(wav_temporary_file)
-        os.remove(mp3_temporary_file)
-
-        waveform = torch.tensor(waveform, dtype=torch.float32)
-        return waveform
-
-
-def main():
-    filename = r"E:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-PH\2025-06-13\active_media_r_2e6e6303-4a2e-4bc9-b814-98ceddc59e9d_23.wav"
-
-    waveform = Mp3ToWavJsonlDataset.filename_to_mp3_waveform(filename, offset=0, duration=15)
-    print(waveform.shape)
-
-    signal = np.array(waveform.numpy() * (1 << 15), dtype=np.int16)
-
-    wavfile.write(
-        "temp.wav",
-        8000,
-        signal,
-    )
-    return
-
-
-if __name__ == "__main__":
-    main()