srt_split.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165
  1. from datetime import datetime
  2. import glob
  3. import os
  4. import soundfile as sf
  5. import re
  6. import tqdm
  7. import pysubs2
  8. ### 부동산 - 전사 srt -> split ogg + 텍스트
  9. root_path = ""
  10. dir_list = glob.glob(root_path + "*/") ## root path 밑에 있는 sub dir 모두 가져오기
  11. pattern_1 = r'\(\(\)\)' ## (())
  12. pattern_2 = r'\{(\w*)\}' ## {@@@} -> @@
  13. pattern_2_1 = r'\[(\w*)\]' ## [@@@] -> @@
  14. pattern_2_2 = r'\((\w*)\)' ## (@@@) 찾기 -> @@
  15. pattern_2_1 = r'[^/]\(([\w|\s]*)\)[^/]' ## (@@)/(@@)이 아닌 (@@) 찾기 repl=' \1 '
  16. pattern_3_1 = r'\(([\w|\s|.]*)\)/\(([\w|\s|.]*)\)' ## (@@)/(@@) -> {@@}/{@@} 임시 변환, repl=r'{\1}/{\2}'
  17. pattern_3_2 = r'\(([\w|\s]*)\)' ## (@@) 찾기, repl=r'\1'
  18. pattern_3_3 = r'\{([\w|\s|.]*)\}/\{([\w|\s|.]*)\}' ## {@@}/{@@} -> (@@)/(@@) 임시변환 해제, r'(\1)/(\2)'
  19. pattern_4 = r'[\(\)\{\}]' ## (, ), {, } 찾기
  20. pattern_4_1 = r'&([\w|-]*)&' ## &(id-id123)& 찾기
  21. pattern_5 = r"[^\uAC00-\uD7A30-9a-zA-Z\s]" ## 특문
  22. pattern_6 = r" {2,}" ## 빈 문자열 2개 이상
  23. pattern_7 = r"\\N"
  24. now = datetime.now()
  25. time_str = now.astimezone().strftime('%Y_%m_%d_%H_%M_%S')
  26. def postprocess(origin):
  27. _a = re.sub(pattern=pattern_7, repl=' ', string = origin)
  28. _a = re.sub(pattern=pattern_1, repl='', string = _a)
  29. _a = re.sub(pattern=pattern_2, repl=r'\1', string = _a)
  30. _a = re.sub(pattern=pattern_2_1, repl=r'\1', string = _a)
  31. # _a = re.sub(pattern=pattern_3, repl=r'\1', string = _a)
  32. _a = re.sub(pattern=pattern_5, repl = "", string = _a)
  33. _a = re.sub(pattern=pattern_6, repl = " ", string = _a)
  34. _a = _a.strip()
  35. if len(_a) == 0:
  36. return None
  37. else:
  38. return _a
  39. def time_convert(time):
  40. split = re.split(r',|:', time)
  41. mi = int(split[0])
  42. mi = mi * 60 + int(split[1])
  43. mi = mi * 60 + int(split[2]) + float('0.' + split[3])
  44. return mi
  45. for _dir in dir_list:
  46. post_none_list = []
  47. not_exist_list = []
  48. ## _dir이 directory일 때만 실행
  49. if not os.path.isdir(_dir):
  50. continue
  51. ## 저장용 split directory 생성
  52. output_dir = os.path.join(_dir, "split")
  53. os.makedirs(output_dir, exist_ok=True)
  54. ## _dir 폴더 하위의 srt 파일 모두 가져오기
  55. srt_list = glob.glob(os.path.join(_dir,'*.srt'))
  56. for srt in tqdm.tqdm(srt_list):
  57. ## 파일명 전처리 -> 폴더 구조별 동작이 다를 수 있으니, 확인 필요
  58. base_name = srt.rsplit('_', 2)[0]
  59. folder_name, file_name = base_name.rsplit('/', 1)
  60. ## srt와 wav가 함께 존재하지 않으면 fail list에 추가 후 continue
  61. if not os.path.isfile(base_name + '.wav'):
  62. not_exist_list.append(file_name)
  63. continue
  64. y, sr = sf.read(base_name + '.wav')
  65. subs = pysubs2.load(srt, encoding="utf-8")
  66. for idx, sub in enumerate(subs):
  67. start = int(sub.start / 1000 * sr)
  68. end = int(sub.end/1000*sr)
  69. ## 전사 내용 후처리 -> 데이터셋별 튜닝 필요.
  70. post = postprocess(sub.text)
  71. ## split file 이름 처리
  72. str_num = str(idx)
  73. str_num = '_' + ('0' * (3 - len(str_num))) + str_num
  74. tmp_file_name = file_name + str_num + ".wav"
  75. ## 전사 내용이 이상값이면 None, None이면 fail list에 추가
  76. if post is None:
  77. post_none_list.append([tmp_file_name, sub.text])
  78. continue
  79. sp_wav = y[start : end+1]
  80. sf.write(os.path.join(output_dir, tmp_file_name), sp_wav, sr)
  81. with open(os.path.join(output_dir, tmp_file_name.replace('wav', 'txt')), 'w+') as f:
  82. f.write(post + '\n')
  83. ## 직접 srt static하게 자르기
  84. # with open(srt, 'r') as f:
  85. # text = ''
  86. # start = 0
  87. # end = 0
  88. # b = False
  89. # for idx, line in enumerate(f):
  90. # if str.isdigit(line):
  91. # b = True
  92. # continue
  93. # if idx % 4 == 1: ## time
  94. # try:
  95. # start, end = map(time_convert, line.split(' --> '))
  96. # except:
  97. # print(srt, '\t', line)
  98. # exit()
  99. # elif idx % 4 == 2: ## transcribe
  100. # str_num = str(idx//4)
  101. # str_num = '_' + ('0' * (3 - len(str_num))) + str_num
  102. # file_name = file_name + str_num + ".wav"
  103. # post = postprocess(line.strip())
  104. # if post is None:
  105. # post_none_list.append([file_name, line])
  106. # continue
  107. # sp_wav = y[int(sr * start) : int(sr * end)+1]
  108. # sf.write(os.path.join(output_dir, file_name), sp_wav, sr)
  109. # with open(os.path.join(output_dir, file_name.replace('wav', 'txt')), 'w+') as f:
  110. # f.write(post + '\n')
  111. if post_none_list:
  112. print('postprocess none = ', len(post_none_list))
  113. txt = f'log/post_none_{time_str}.txt'
  114. with open(txt, 'w+') as f:
  115. for item in post_none_list:
  116. f.write("\t".join(item) + "\n")
  117. if not_exist_list:
  118. print('not exist = ', len(not_exist_list))
  119. txt = f'log/not_exist_{time_str}.txt'
  120. for item in post_none_list:
  121. with open(txt, 'w+') as f:
  122. f.write("\t".join(item) + "\n")