| 1 | import json |
| 2 | import os.path |
| 3 | import re |
| 4 | from timeit import default_timer as timer |
| 5 | |
| 6 | try: |
| 7 | from faster_whisper import WhisperModel |
| 8 | except ImportError: |
| 9 | WhisperModel = None |
| 10 | from loguru import logger |
| 11 | |
| 12 | from app.config import config |
| 13 | from app.utils import utils |
| 14 | |
| 15 | model_size = config.whisper.get("model_size", "large-v3") |
| 16 | device = config.whisper.get("device", "cpu") |
| 17 | compute_type = config.whisper.get("compute_type", "int8") |
| 18 | model = None |
| 19 | |
| 20 | |
| 21 | def create(audio_file, subtitle_file: str = ""): |
| 22 | global model |
| 23 | if WhisperModel is None: |
| 24 | logger.warning("faster_whisper not available, skipping whisper subtitle generation") |
| 25 | return "" |
| 26 | if not model: |
| 27 | model_path = f"{utils.root_dir()}/models/whisper-{model_size}" |
| 28 | model_bin_file = f"{model_path}/model.bin" |
| 29 | if not os.path.isdir(model_path) or not os.path.isfile(model_bin_file): |
| 30 | model_path = model_size |
| 31 | |
| 32 | logger.info( |
| 33 | f"loading model: {model_path}, device: {device}, compute_type: {compute_type}" |
| 34 | ) |
| 35 | try: |
| 36 | model = WhisperModel( |
| 37 | model_size_or_path=model_path, device=device, compute_type=compute_type |
| 38 | ) |
| 39 | except Exception as e: |
| 40 | logger.error( |
| 41 | f"failed to load model: {e} \n\n" |
| 42 | f"********************************************\n" |
| 43 | f"this may be caused by network issue. \n" |
| 44 | f"please download the model manually and put it in the 'models' folder. \n" |
| 45 | f"see [README.md FAQ](https://github.com/harry0703/MoneyPrinterTurbo) for more details.\n" |
| 46 | f"********************************************\n\n" |
| 47 | ) |
| 48 | return None |
| 49 | |
| 50 | logger.info(f"start, output file: {subtitle_file}") |
| 51 | if not subtitle_file: |
| 52 | subtitle_file = f"{audio_file}.srt" |
| 53 | |
| 54 | segments, info = model.transcribe( |
| 55 | audio_file, |
| 56 | beam_size=5, |
| 57 | word_timestamps=True, |
| 58 | vad_filter=True, |
| 59 | vad_parameters=dict(min_silence_duration_ms=500), |
| 60 | ) |
| 61 | |
| 62 | logger.info( |
| 63 | f"detected language: '{info.language}', probability: {info.language_probability:.2f}" |
| 64 | ) |
| 65 | |
| 66 | start = timer() |
| 67 | subtitles = [] |
| 68 | |
| 69 | def recognized(seg_text, seg_start, seg_end): |
| 70 | seg_text = seg_text.strip() |
| 71 | if not seg_text: |
| 72 | return |
| 73 | |
| 74 | msg = "[%.2fs -> %.2fs] %s" % (seg_start, seg_end, seg_text) |
| 75 | logger.debug(msg) |
| 76 | |
| 77 | subtitles.append( |
| 78 | {"msg": seg_text, "start_time": seg_start, "end_time": seg_end} |
| 79 | ) |
| 80 | |
| 81 | for segment in segments: |
| 82 | words_idx = 0 |
| 83 | words_len = len(segment.words) |
| 84 | |
| 85 | seg_start = 0 |
| 86 | seg_end = 0 |
| 87 | seg_text = "" |
| 88 | |
| 89 | if segment.words: |
| 90 | is_segmented = False |
| 91 | for word in segment.words: |
| 92 | if not is_segmented: |
| 93 | seg_start = word.start |
| 94 | is_segmented = True |
| 95 | |
| 96 | seg_end = word.end |
| 97 | # If it contains punctuation, then break the sentence. |
| 98 | seg_text += word.word |
| 99 | |
| 100 | if utils.str_contains_punctuation(word.word): |
| 101 | # remove last char |
| 102 | seg_text = seg_text[:-1] |
| 103 | if not seg_text: |
| 104 | continue |
| 105 | |
| 106 | recognized(seg_text, seg_start, seg_end) |
| 107 | |
| 108 | is_segmented = False |
| 109 | seg_text = "" |
| 110 | |
| 111 | if words_idx == 0 and segment.start < word.start: |
| 112 | seg_start = word.start |
| 113 | if words_idx == (words_len - 1) and segment.end > word.end: |
| 114 | seg_end = word.end |
| 115 | words_idx += 1 |
| 116 | |
| 117 | if not seg_text: |
| 118 | continue |
| 119 | |
| 120 | recognized(seg_text, seg_start, seg_end) |
| 121 | |
| 122 | end = timer() |
| 123 | |
| 124 | diff = end - start |
| 125 | logger.info(f"complete, elapsed: {diff:.2f} s") |
| 126 | |
| 127 | idx = 1 |
| 128 | lines = [] |
| 129 | for subtitle in subtitles: |
| 130 | text = subtitle.get("msg") |
| 131 | if text: |
| 132 | lines.append( |
| 133 | utils.text_to_srt( |
| 134 | idx, text, subtitle.get("start_time"), subtitle.get("end_time") |
| 135 | ) |
| 136 | ) |
| 137 | idx += 1 |
| 138 | |
| 139 | sub = "\n".join(lines) + "\n" |
| 140 | with open(subtitle_file, "w", encoding="utf-8") as f: |
| 141 | f.write(sub) |
| 142 | logger.info(f"subtitle file created: {subtitle_file}") |
| 143 | |
| 144 | |
| 145 | def file_to_subtitles(filename): |
| 146 | if not filename or not os.path.isfile(filename): |
| 147 | return [] |
| 148 | |
| 149 | times_texts = [] |
| 150 | current_times = None |
| 151 | current_text = "" |
| 152 | index = 0 |
| 153 | with open(filename, "r", encoding="utf-8") as f: |
| 154 | for line in f: |
| 155 | times = re.findall("([0-9]*:[0-9]*:[0-9]*,[0-9]*)", line) |
| 156 | if times: |
| 157 | current_times = line |
| 158 | elif line.strip() == "" and current_times: |
| 159 | index += 1 |
| 160 | times_texts.append((index, current_times.strip(), current_text.strip())) |
| 161 | current_times, current_text = None, "" |
| 162 | elif current_times: |
| 163 | current_text += line |
| 164 | |
| 165 | # Flush the final block. SRT files whose last subtitle is not followed by a |
| 166 | # trailing blank line never hit the blank-line branch above, so without this |
| 167 | # the last subtitle would be silently dropped. |
| 168 | if current_times: |
| 169 | index += 1 |
| 170 | times_texts.append((index, current_times.strip(), current_text.strip())) |
| 171 | return times_texts |
| 172 | |
| 173 | |
| 174 | def levenshtein_distance(s1, s2): |
| 175 | if len(s1) < len(s2): |
| 176 | return levenshtein_distance(s2, s1) |
| 177 | |
| 178 | if len(s2) == 0: |
| 179 | return len(s1) |
| 180 | |
| 181 | previous_row = range(len(s2) + 1) |
| 182 | for i, c1 in enumerate(s1): |
| 183 | current_row = [i + 1] |
| 184 | for j, c2 in enumerate(s2): |
| 185 | insertions = previous_row[j + 1] + 1 |
| 186 | deletions = current_row[j] + 1 |
| 187 | substitutions = previous_row[j] + (c1 != c2) |
| 188 | current_row.append(min(insertions, deletions, substitutions)) |
| 189 | previous_row = current_row |
| 190 | |
| 191 | return previous_row[-1] |
| 192 | |
| 193 | |
| 194 | def similarity(a, b): |
| 195 | distance = levenshtein_distance(a.lower(), b.lower()) |
| 196 | max_length = max(len(a), len(b)) |
| 197 | return 1 - (distance / max_length) |
| 198 | |
| 199 | |
| 200 | def correct(subtitle_file, video_script): |
| 201 | subtitle_items = file_to_subtitles(subtitle_file) |
| 202 | normalized_script = utils.normalize_script_for_subtitle_matching(video_script) |
| 203 | script_lines = utils.split_string_by_punctuations(normalized_script) |
| 204 | |
| 205 | corrected = False |
| 206 | new_subtitle_items = [] |
| 207 | script_index = 0 |
| 208 | subtitle_index = 0 |
| 209 | |
| 210 | while script_index < len(script_lines) and subtitle_index < len(subtitle_items): |
| 211 | script_line = script_lines[script_index].strip() |
| 212 | subtitle_line = subtitle_items[subtitle_index][2].strip() |
| 213 | |
| 214 | if script_line == subtitle_line: |
| 215 | new_subtitle_items.append(subtitle_items[subtitle_index]) |
| 216 | script_index += 1 |
| 217 | subtitle_index += 1 |
| 218 | else: |
| 219 | combined_subtitle = subtitle_line |
| 220 | start_time = subtitle_items[subtitle_index][1].split(" --> ")[0] |
| 221 | end_time = subtitle_items[subtitle_index][1].split(" --> ")[1] |
| 222 | next_subtitle_index = subtitle_index + 1 |
| 223 | |
| 224 | while next_subtitle_index < len(subtitle_items): |
| 225 | next_subtitle = subtitle_items[next_subtitle_index][2].strip() |
| 226 | if similarity( |
| 227 | script_line, combined_subtitle + " " + next_subtitle |
| 228 | ) > similarity(script_line, combined_subtitle): |
| 229 | combined_subtitle += " " + next_subtitle |
| 230 | end_time = subtitle_items[next_subtitle_index][1].split(" --> ")[1] |
| 231 | next_subtitle_index += 1 |
| 232 | else: |
| 233 | break |
| 234 | |
| 235 | if similarity(script_line, combined_subtitle) > 0.8: |
| 236 | logger.warning( |
| 237 | f"Merged/Corrected - Script: {script_line}, Subtitle: {combined_subtitle}" |
| 238 | ) |
| 239 | new_subtitle_items.append( |
| 240 | ( |
| 241 | len(new_subtitle_items) + 1, |
| 242 | f"{start_time} --> {end_time}", |
| 243 | script_line, |
| 244 | ) |
| 245 | ) |
| 246 | corrected = True |
| 247 | else: |
| 248 | logger.warning( |
| 249 | f"Mismatch - Script: {script_line}, Subtitle: {combined_subtitle}" |
| 250 | ) |
| 251 | new_subtitle_items.append( |
| 252 | ( |
| 253 | len(new_subtitle_items) + 1, |
| 254 | f"{start_time} --> {end_time}", |
| 255 | script_line, |
| 256 | ) |
| 257 | ) |
| 258 | corrected = True |
| 259 | |
| 260 | script_index += 1 |
| 261 | subtitle_index = next_subtitle_index |
| 262 | |
| 263 | # Process the remaining lines of the script. |
| 264 | while script_index < len(script_lines): |
| 265 | logger.warning(f"Extra script line: {script_lines[script_index]}") |
| 266 | if subtitle_index < len(subtitle_items): |
| 267 | new_subtitle_items.append( |
| 268 | ( |
| 269 | len(new_subtitle_items) + 1, |
| 270 | subtitle_items[subtitle_index][1], |
| 271 | script_lines[script_index], |
| 272 | ) |
| 273 | ) |
| 274 | subtitle_index += 1 |
| 275 | else: |
| 276 | new_subtitle_items.append( |
| 277 | ( |
| 278 | len(new_subtitle_items) + 1, |
| 279 | "00:00:00,000 --> 00:00:00,000", |
| 280 | script_lines[script_index], |
| 281 | ) |
| 282 | ) |
| 283 | script_index += 1 |
| 284 | corrected = True |
| 285 | |
| 286 | if corrected: |
| 287 | with open(subtitle_file, "w", encoding="utf-8") as fd: |
| 288 | for i, item in enumerate(new_subtitle_items): |
| 289 | fd.write(f"{i + 1}\n{item[1]}\n{item[2]}\n\n") |
| 290 | logger.info("Subtitle corrected") |
| 291 | else: |
| 292 | logger.success("Subtitle is correct") |
| 293 | |
| 294 | |
| 295 | if __name__ == "__main__": |
| 296 | task_id = "c12fd1e6-4b0a-4d65-a075-c87abe35a072" |
| 297 | task_dir = utils.task_dir(task_id) |
| 298 | subtitle_file = f"{task_dir}/subtitle.srt" |
| 299 | audio_file = f"{task_dir}/audio.mp3" |
| 300 | |
| 301 | subtitles = file_to_subtitles(subtitle_file) |
| 302 | print(subtitles) |
| 303 | |
| 304 | script_file = f"{task_dir}/script.json" |
| 305 | with open(script_file, "r") as f: |
| 306 | script_content = f.read() |
| 307 | s = json.loads(script_content) |
| 308 | script = s.get("script") |
| 309 | |
| 310 | correct(subtitle_file, script) |
| 311 | |
| 312 | subtitle_file = f"{task_dir}/subtitle-test.srt" |
| 313 | create(audio_file, subtitle_file) |
| 314 |