from preprocess.NAT_mel import MelNet import os from tqdm import tqdm from glob import glob import math import pandas as pd import argparse from argparse import Namespace import math import audioread from tqdm.contrib.concurrent import process_map import torch import torch.nn as nn import torchaudio import numpy as np from torch.distributed import init_process_group from torch.utils.data import Dataset,DataLoader,DistributedSampler import torch.multiprocessing as mp import json class tsv_dataset(Dataset): def __init__(self,tsv_path,sr,mode='none',hop_size = None,target_mel_length = None) -> None: super().__init__() if os.path.isdir(tsv_path): files = glob(os.path.join(tsv_path,'*.tsv')) df = pd.concat([pd.read_csv(file,sep='\t') for file in files]) else: df = pd.read_csv(tsv_path,sep='\t') self.audio_paths = [] self.sr = sr self.mode = mode self.target_mel_length = target_mel_length self.hop_size = hop_size for t in tqdm(df.itertuples()): self.audio_paths.append(getattr(t,'audio_path')) def __len__(self): return len(self.audio_paths) def pad_wav(self,wav): # wav should be in shape(1,wav_len) wav_length = wav.shape[-1] assert wav_length > 100, "wav is too short, %s" % wav_length segment_length = (self.target_mel_length + 1) * self.hop_size # final mel will crop the last mel, mel = mel[:,:-1] if segment_length is None or wav_length == segment_length: return wav elif wav_length > segment_length: return wav[:,:segment_length] elif wav_length < segment_length: temp_wav = torch.zeros((1, segment_length),dtype=torch.float32) temp_wav[:, :wav_length] = wav return temp_wav def __getitem__(self, index): audio_path = self.audio_paths[index] wav, orisr = torchaudio.load(audio_path) if wav.shape[0] != 1: # stereo to mono (2,wav_len) -> (1,wav_len) wav = wav.mean(0,keepdim=True) wav = torchaudio.functional.resample(wav, orig_freq=orisr, new_freq=self.sr) if self.mode == 'pad': assert self.target_mel_length is not None wav = self.pad_wav(wav) return audio_path,wav def process_audio_by_tsv(rank,args): if args.num_gpus > 1: init_process_group(backend=args.dist_config['dist_backend'], init_method=args.dist_config['dist_url'], world_size=args.dist_config['world_size'] * args.num_gpus, rank=rank) sr = args.audio_sample_rate dataset = tsv_dataset(args.tsv_path,sr = sr,mode=args.mode,hop_size=args.hop_size,target_mel_length=args.batch_max_length) sampler = DistributedSampler(dataset,shuffle=False) if args.num_gpus > 1 else None # batch_size must == 1,since wav_len is not equal loader = DataLoader(dataset, sampler=sampler,batch_size=1, num_workers=16,drop_last=False) device = torch.device('cuda:{:d}'.format(rank)) mel_net = MelNet(args.__dict__) mel_net.to(device) # if args.num_gpus > 1: # RuntimeError: DistributedDataParallel is not needed when a module doesn't have any parameter that requires a gradient. # mel_net = DistributedDataParallel(mel_net, device_ids=[rank]).to(device) root = args.save_path loader = tqdm(loader) if rank == 0 else loader for batch in loader: audio_paths,wavs = batch wavs = wavs.to(device) if args.save_resample: for audio_path,wav in zip(audio_paths,wavs): psplits = audio_path.split('/') wav_name = psplits[-1] # save resample resample_root,resample_name = root+f'_{sr}',wav_name[:-4]+'_audio.npy' resample_dir_name = os.path.join(resample_root,*psplits[1:-1]) resample_path = os.path.join(resample_dir_name,resample_name) os.makedirs(resample_dir_name,exist_ok=True) np.save(resample_path,wav.cpu().numpy().squeeze(0)) if args.save_mel: mode = args.mode batch_max_length = args.batch_max_length for audio_path,wav in zip(audio_paths,wavs): psplits = audio_path.split('/') wav_name = psplits[-1] mel_root,mel_name = root,wav_name[:-4]+'_mel.npy' mel_dir_name = os.path.join(mel_root,f'mel{mode}{sr}',*psplits[1:-1]) mel_path = os.path.join(mel_dir_name,mel_name) if not os.path.exists(mel_path): mel_spec = mel_net(wav).cpu().numpy().squeeze(0) # (mel_bins,mel_len) if mel_spec.shape[1] <= batch_max_length: if mode == 'tile': # pad is done in dataset as pad wav n_repeat = math.ceil((batch_max_length + 1) / mel_spec.shape[1]) mel_spec = np.tile(mel_spec,reps=(1,n_repeat)) elif mode == 'none' or mode == 'pad': pass else: raise ValueError(f'mode:{mode} is not supported') mel_spec = mel_spec[:,:batch_max_length] os.makedirs(mel_dir_name,exist_ok=True) np.save(mel_path,mel_spec) def split_list(i_list,num): each_num = math.ceil(i_list / num) result = [] for i in range(num): s = each_num * i e = (each_num * (i+1)) result.append(i_list[s:e]) return result def drop_bad_wav(item): index,path = item try: with audioread.audio_open(path) as f: totalsec = f.duration if totalsec < 0.1: return index # index except: print(f"corrupted wav:{path}") return index return False def drop_bad_wavs(tsv_path):# 'audioset.csv' df = pd.read_csv(tsv_path,sep='\t') item_list = [] for item in tqdm(df.itertuples()): item_list.append((item[0],getattr(item,'audio_path'))) r = process_map(drop_bad_wav,item_list,max_workers=16,chunksize=16) bad_indices = list(filter(lambda x:x!= False,r)) print(bad_indices) with open('bad_wavs.json','w') as f: x = [item_list[i] for i in bad_indices] json.dump(x,f) df = df.drop(bad_indices,axis=0) df.to_csv(tsv_path,sep='\t',index=False) def addmel2tsv(save_dir,tsv_path): df = pd.read_csv(tsv_path,sep='\t') mels = glob(f'{save_dir}/mel{args.mode}{args.audio_sample_rate}/**/*_mel.npy',recursive=True) name2mel,idx2name,idx2mel = {},{},{} for mel in mels: bn = os.path.basename(mel)[:-8]# remove _mel.npy name2mel[bn] = mel for t in df.itertuples(): idx = int(t[0]) bn = os.path.basename(getattr(t,'audio_path'))[:-4] idx2name[idx] = bn for k,v in idx2name.items(): idx2mel[k] = name2mel[v] df['mel_path'] = df.index.map(idx2mel) df.to_csv(tsv_path,sep='\t',index=False) def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "--tsv_path",type=str) parser.add_argument( "--num_gpus",type=int,default=1) parser.add_argument( "--max_duration",type=int,default=30) return parser.parse_args() if __name__ == '__main__': pargs = parse_args() tsv_path = pargs.tsv_path if os.path.isdir(tsv_path): files = glob(os.path.join(tsv_path,'*.tsv')) for file in files: drop_bad_wavs(file) else: drop_bad_wavs(tsv_path) num_gpus = pargs.num_gpus batch_max_length = int(pargs.max_duration * 62.5)# 62.5 is the mel length for 1 second save_path = 'processed' args = { 'audio_sample_rate': 16000, 'audio_num_mel_bins':80, 'fft_size': 1024, 'win_size': 1024, 'hop_size': 256, 'fmin': 0, 'fmax': 8000, 'batch_max_length': batch_max_length, 'tsv_path': tsv_path, 'num_gpus': num_gpus, 'mode': 'none', # pad,none, 'save_resample':False, 'save_mel' :True, 'save_path': save_path, } os.makedirs(save_path,exist_ok=True) args = Namespace(**args) args.dist_config = { "dist_backend": "nccl", "dist_url": "tcp://localhost:54189", "world_size": 1 } if args.num_gpus>1: mp.spawn(process_audio_by_tsv,nprocs=args.num_gpus,args=(args,)) else: process_audio_by_tsv(0,args=args) print("proceoss mel done") addmel2tsv(save_path,tsv_path) print("done")