transcribe_process.py 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232
  1. import jiwer
  2. import glob
  3. import pandas as pd
  4. import tqdm
  5. import os
  6. import re
  7. from collections import defaultdict
  8. import matplotlib.pyplot as plt
  9. root_path = ""
  10. sub_path = [
  11. "", ""
  12. ]
  13. def readNumber(n):
  14. units = [''] + list('십백천만') + ['십만', '백만', '천만', '억']
  15. # units = [''] + list('십백천만')
  16. nums = '일이삼사오육칠팔구'
  17. result = []
  18. i = 0
  19. while n > 0:
  20. n, r = divmod(n, 10)
  21. if r > 0:
  22. result.append(nums[r-1] + units[i])
  23. i += 1
  24. return ''.join(result[::-1])
  25. def convert(str):
  26. iter = re.finditer(r'\d+', str)
  27. L = sum(1 for _ in re.finditer(r'\d+', str))
  28. res = ''
  29. prev = 0
  30. if L == 0:
  31. return str
  32. for idx, match in enumerate(iter):
  33. s, e = match.span()
  34. if prev < s:
  35. res += str[prev:s]
  36. a = str[s : e]
  37. a = readNumber(int(a))
  38. res += a
  39. prev = e
  40. if idx == L - 1:
  41. res += str[e:]
  42. return res
  43. def create_total(total_name):
  44. for sub in sub_path:
  45. srt_list = glob.glob(root_path + sub + "*.srt")
  46. with open(root_path + f"total/{total_name}", 'a') as tt:
  47. for srt in srt_list:
  48. base_name = srt.rsplit('_', 2)[0]
  49. folder_name, file_name = base_name.rsplit('/', 1)
  50. with open(srt, 'r') as f:
  51. split_text = ' '
  52. for idx, org_text in enumerate(f):
  53. if idx%4 == 2:
  54. str_num = str(idx//4)
  55. str_num = '_' + ('0' * (3 - len(str_num))) + str_num
  56. split_file = folder_name + "/split/" + file_name + str_num + ".txt"
  57. if not os.path.exists(split_file):
  58. split_text = ''
  59. else:
  60. with open(split_file, 'r') as f:
  61. split_text = f.read()
  62. org_text = re.sub(r"[^\uAC00-\uD7A30-9a-zA-Z\s]", "", org_text.strip())
  63. split_text = re.sub(r"[^\uAC00-\uD7A30-9a-zA-Z\s]", "", split_text.strip())
  64. # org_text = re.sub(r"[^\uAC00-\uD7A30-9a-zA-Z]", "", org_text.strip())
  65. # split_text = re.sub(r"[^\uAC00-\uD7A30-9a-zA-Z]", "", split_text.strip())
  66. if len(org_text) == 0 or len(split_text) == 0:
  67. cer, wer, c_cer, c_wer = [ 1 for _ in range(4) ]
  68. else:
  69. cer = jiwer.cer(org_text, split_text)
  70. wer = jiwer.wer(org_text, split_text)
  71. try:
  72. c_org = convert(org_text)
  73. c_split = convert(split_text)
  74. c_cer = jiwer.cer(c_org, c_split)
  75. c_wer = jiwer.wer(c_org, c_split)
  76. except:
  77. c_cer = cer
  78. c_wer = wer
  79. tt.write(f"{file_name + str_num}\t{org_text}\t{split_text}\t{cer}\t{wer}\t{c_cer}\t{c_wer}\n")
  80. def create_dataframe(total_name):
  81. with open(root_path + f"total/{total_name}", 'r') as f:
  82. data = f.readlines()
  83. columns = ['path', 'script1', 'script2', 'cer', 'wer', 'c_cer', 'c_wer']
  84. data = [ d.strip().split('\t') for d in data ]
  85. df = pd.DataFrame(data, columns=columns)
  86. threshold = '0.2'
  87. print(df[(df['cer'] < threshold) | (df['wer'] < threshold) | (df['c_cer'] < threshold) | (df['c_wer'] < threshold) ][['path', 'cer', 'wer', 'c_cer', 'c_wer']])
  88. print('-'*50)
  89. print(df[(df['cer'] < threshold) | (df['wer'] < threshold) | (df['c_cer'] < threshold) | (df['c_wer'] < threshold) ][['path', 'script1', 'script2']])
  90. print('-'*50)
  91. threshold = '0.0'
  92. print(df[(df['cer'] <= threshold) | (df['wer'] <= threshold) | (df['c_cer'] <= threshold) | (df['c_wer'] <= threshold) ][['path', 'script1', 'script2']])
  93. print('-'*50)
  94. def create_plt_file(d, name):
  95. plt.title(name)
  96. plt.xlim([0, 3])
  97. plt.bar(d.keys(), d.values(), 0.01, color='g')
  98. plt.savefig(f'plot/{name}.png')
  99. def create_hist():
  100. hist_cer = defaultdict(int)
  101. hist_wer = defaultdict(int)
  102. hist_c_cer = defaultdict(int)
  103. hist_c_wer = defaultdict(int)
  104. error_cnt = 0
  105. for sub in sub_path:
  106. srt_list = glob.glob(root_path + sub + "*.srt")
  107. for srt in srt_list:
  108. base_name = srt.rsplit('_', 2)[0]
  109. folder_name, file_name = base_name.rsplit('/', 1)
  110. with open(srt, 'r') as f:
  111. split_text = ' '
  112. for idx, org_text in enumerate(f):
  113. cer, wer, c_cer, c_wer = [ 0 for _ in range(4) ]
  114. if idx%4 == 2:
  115. str_num = str(idx//4)
  116. str_num = '_' + ('0' * (3 - len(str_num))) + str_num
  117. split_file = folder_name + "/split/" + file_name + str_num + ".txt"
  118. if not os.path.exists(split_file):
  119. split_text = ''
  120. else:
  121. with open(split_file, 'r') as f:
  122. split_text = f.read()
  123. org_text = org_text.strip()
  124. split_text = split_text.strip()
  125. if len(org_text) == 0 or len(split_text) == 0:
  126. e = 1.0
  127. hist_cer[e] += 1
  128. hist_wer[e] += 1
  129. hist_c_cer[e] += 1
  130. hist_c_wer[e] += 1
  131. else:
  132. cer = round(jiwer.cer(org_text, split_text), 4)
  133. wer = round(jiwer.wer(org_text, split_text), 4)
  134. hist_cer[cer] += 1
  135. hist_wer[wer] += 1
  136. try:
  137. c_org = convert(org_text)
  138. c_split = convert(split_text)
  139. c_cer = round(jiwer.cer(c_org, c_split), 4)
  140. c_wer = round(jiwer.wer(c_org, c_split), 4)
  141. hist_c_cer[c_cer] += 1
  142. hist_c_wer[c_wer] += 1
  143. except:
  144. error_cnt += 1
  145. pass
  146. # create_plt_file(hist_cer, 'cer')
  147. # create_plt_file(hist_wer, 'wer')
  148. # create_plt_file(hist_c_cer, 'convert cer')
  149. # create_plt_file(hist_c_wer, 'convert wer')
  150. # print(error_cnt)
  151. width = 0.005
  152. xlim = [0, 1.1]
  153. # plt.subplot(2, 2, 1)
  154. # plt.title('cer')
  155. # plt.xlim(xlim)
  156. # plt.bar(hist_cer.keys(), hist_cer.values(), width, color='g')
  157. # plt.subplot(2, 2, 2)
  158. # plt.title('wer')
  159. # plt.xlim(xlim)
  160. # plt.bar(hist_wer.keys(), hist_wer.values(), width, color='g')
  161. # plt.subplot(2, 2, 3)
  162. # plt.title('convert cer')
  163. # plt.xlim(xlim)
  164. # plt.bar(hist_c_cer.keys(), hist_c_cer.values(), width, color='g')
  165. # plt.subplot(2, 2, 4)
  166. # plt.title('convert wer')
  167. # plt.xlim(xlim)
  168. # plt.bar(hist_c_wer.keys(), hist_c_wer.values(), width, color='g')
  169. # plt.tight_layout()
  170. # plt.savefig(f'plot/total.png', dpi=200)
  171. if __name__ == '__main__':
  172. total_name = "total_without_space.txt"
  173. # create_hist()
  174. create_total(total_name)
  175. create_dataframe(total_name)