check_ogg.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173
  1. import argparse
  2. import numpy as np
  3. import random
  4. import os
  5. import time
  6. import math
  7. import glob
  8. import soundfile as sf
  9. from scipy import signal
  10. from scipy.io import wavfile
  11. import librosa
  12. import io
  13. import torch
  14. # import torchaudio
  15. import torchaudio.transforms as T
  16. import torch.nn.functional as F
  17. from tqdm import tqdm
  18. import multiprocessing as mp
  19. from multiprocessing import Pool
  20. from typing import Any, Coroutine, Iterable, List, Tuple
  21. def loadWav_samplerate(filepath, sample_rate, transform=None):
  22. waveform, _sample_rate = sf.read(filepath)
  23. if _sample_rate != sample_rate:
  24. t = torch.from_numpy(waveform).type(torch.float32)
  25. if transform is None:
  26. transform = T.Resample(_sample_rate, sample_rate)
  27. waveform = transform(t).numpy()
  28. return waveform
  29. def check_ogg(root_dir, wav_path, sample_rate = 16000):
  30. try:
  31. sub_path = wav_path.rsplit('.', 1)[0].split(root_dir)[-1] + '.ogg' ## wav_path = root_dir + sub_path
  32. if sub_path[0] == '/':
  33. sub_path = sub_path[1:]
  34. ogg_path = os.path.join(root_dir, 'ogg', sub_path)
  35. errors = []
  36. if not os.path.isfile(ogg_path):
  37. error = {'status' : "ogg file not exist", 'detail' : ogg_path}
  38. errors.append(error)
  39. return {'status' : 'fail', 'file' : wav_path, 'errors' : errors}
  40. wav_sound = loadWav_samplerate(wav_path, sample_rate)
  41. file_format = "OGG"
  42. memory_file = io.BytesIO( )
  43. memory_file.name = "test.ogg"
  44. sf.write( memory_file, wav_sound, sample_rate, format = file_format )
  45. memory_file.seek( 0 )
  46. temp_data, temp_sr = sf.read( memory_file )
  47. ogg_sound, sr = sf.read(ogg_path)
  48. if sr != sample_rate:
  49. error = {'status' : "sample rate not match", 'detail' : str(sr) + "\t" + str(sample_rate)}
  50. errors.append(error)
  51. if not np.array_equal(temp_data, ogg_sound):
  52. error = {'status' : "wav and ogg not equal"}
  53. errors.append(error)
  54. return {'status' : 'fail', 'file' : wav_path, 'errors' : errors}
  55. except Exception as e:
  56. return {'status' : 'exception', 'file' : wav_path, 'errors' : [str(e)]}
  57. return {'status' : 'success', 'file' : wav_path}
  58. def str2bool(v):
  59. if isinstance(v, bool):
  60. return v
  61. if v.lower() in ('yes', 'true', 't', 'y', '1'):
  62. return True
  63. elif v.lower() in ('no', 'false', 'f', 'n', '0'):
  64. return False
  65. else:
  66. raise argparse.ArgumentTypeError('Boolean value expected.')
  67. if __name__ == '__main__':
  68. parser = argparse.ArgumentParser()
  69. parser.add_argument('--root-dir', default='/root/nas/data/', type=str, help='setting to the root directory')
  70. parser.add_argument('--sub-dir', default='clovacall', type=str, help='setting to the root directory')
  71. parser.add_argument('--dest-dir', default='/root/preprocess/backup', type=str, help='The destination save path of ogg audio file')
  72. parser.add_argument('--use-list', default=False, type=str2bool, help='use a files.txt')
  73. args = parser.parse_args()
  74. print("task dir : ", os.path.join(args.root_dir, args.sub_dir))
  75. wav_files = []
  76. files_txt = args.sub_dir + "_wavs.txt"
  77. if args.use_list and os.path.exists(files_txt):
  78. with open(files_txt, 'r') as f:
  79. for line in f:
  80. wav_files.append(line.split('\n')[0])
  81. else:
  82. print("glob start !!")
  83. wav_files = glob.glob(os.path.join(args.root_dir, args.sub_dir, '**/*.wav'), recursive=True)
  84. print("glob end !!")
  85. if args.use_list and not os.path.isfile(files_txt):
  86. with open(files_txt, 'w+') as f:
  87. for wav in wav_files:
  88. f.write(f'{wav}\n')
  89. # wav_files = wav_files[:100] ## for test
  90. # for wav_path in wav_files:
  91. # r = os.path.join(args.root_dir, args.sub_dir)
  92. # sub_path = wav_path.rsplit('.', 1)[0].split(r)[-1] + '.ogg' ## wav_path = root_dir + sub_path
  93. # if sub_path[0] == '/':
  94. # sub_path = sub_path[1:]
  95. # ogg_path = os.path.join(r, 'ogg', sub_path)
  96. # # ogg_path = f"{os.path.join(args.root_dir, args.sub_dir)}/ogg/{sub_path}"
  97. # print(sub_path + "\t" + ogg_path + "\t" + r)
  98. print("wav len = ", len(wav_files))
  99. # exit()
  100. num_cores = mp.cpu_count()
  101. def work_func(path):
  102. return check_ogg(os.path.join(args.root_dir, args.sub_dir), path, 16000)
  103. fail_list = []
  104. try:
  105. pool = Pool((num_cores * 2) // 3)
  106. with tqdm(total=len(wav_files)) as pbar:
  107. for res in tqdm(pool.imap_unordered(work_func, wav_files)):
  108. # print(res)
  109. if res['status'] != 'success':
  110. fail_list.append(res)
  111. pbar.update()
  112. pool.close()
  113. pool.join()
  114. fail_txt = "fail_" + args.sub_dir + ".txt"
  115. if len(fail_list) >= 1:
  116. with open(fail_txt, 'w+') as f:
  117. for fail in fail_list:
  118. content = res["file"] + "\t" + str(res['errors'])
  119. f.write(content + "\n")
  120. print("Fail len ", len(fail_list))
  121. print("Script complete!")
  122. except KeyboardInterrupt:
  123. pool.close()
  124. pool.join()