使用 Whisper 生成带时间戳的 SRT 字幕文件及进阶应用

[
  {
    "id": 0,
    "seek": 0,
    "start": 0.0,
    "end": 3.4,
    "text": " 这是一个示例。",
    "tokens": [50363, 261, 264, 5322, 1079, 286, 264, 1374, 13, 50533],
    "temperature": 0.0,
    "avg_logprob": -0.21,
    "compression_ratio": 1.2,
    "no_speech_prob": 0.001
  },
  {
    "id": 1,
    "seek": 3.4,
    "start": 3.4,
    "end": 6.8,
    "text": " 我们将学习如何生成SRT文件。",
    "tokens": [50533, 261, 264, 5322, 1079, 286, 264, 1374, 13, 50533],
    "temperature": 0.0,
    "avg_logprob": -0.21,
    "compression_ratio": 1.2,
    "no_speech_prob": 0.001
  }
  // ... 更多片段
]

1
00:00:00,000 --> 00:00:03,400
这是一个示例。

2
00:00:03,400 --> 00:00:06,800
我们将学习如何生成SRT文件。

import os
import whisper
from tqdm import tqdm
import datetime

# 定义存储mp4文件的根目录
root_folder = "C:\Video"

# 加载 Whisper 模型
print("Loading whisper model...")
model = whisper.load_model("base") # 可以根据需求选择 "tiny", "base", "small", "medium", "large"
print("Whisper model complete.")

# 获取待转录的mp4文件数量
print("Getting number of files to transcribe...")
num_files = sum(1 for dirpath, dirnames, filenames in os.walk(root_folder) for filename in filenames if filename.endswith(".mp4"))
print("Number of files: ", num_files)

# 转录mp4文件并显示进度条
with tqdm(total=num_files, desc="Transcribing Files") as pbar:
    for dirpath, dirnames, filenames in os.walk(root_folder):
        for filename in filenames:
            if filename.endswith(".mp4"):
                filepath = os.path.join(dirpath, filename)
                print(f"\nTranscribing: {filename}")
                # 进行转录，fp16=False 适用于没有GPU或GPU不支持FP16的情况
                result = model.transcribe(filepath, fp16=False, verbose=True)

                # 获取不带扩展名的文件名
                filename_no_ext = os.path.splitext(filename)[0]

                # 生成 SRT 文件
                srt_filepath = os.path.join(dirpath, filename_no_ext + '.srt')
                with open(srt_filepath, 'w', encoding='utf-8') as f_srt:
                    for i, segment in enumerate(result['segments']):
                        # SRT 片段序号
                        f_srt.write(str(i + 1) + '\n')

                        # 格式化时间戳
                        start_time = str(datetime.timedelta(seconds=segment['start']))
                        end_time = str(datetime.timedelta(seconds=segment['end']))

                        # 处理毫秒部分，timedelta 默认是微秒，需要转换
                        # 格式化为 HH:MM:SS,msmsms
                        start_ms = int((segment['start'] % 1) * 1000)
                        end_ms = int((segment['end'] % 1) * 1000)

                        start_str = f"{int(segment['start'] // 3600):02}:{int((segment['start'] % 3600) // 60):02}:{int(segment['start'] % 60):02},{start_ms:03}"
                        end_str = f"{int(segment['end'] // 3600):02}:{int((segment['end'] % 3600) // 60):02}:{int(segment['end'] % 60):02},{end_ms:03}"

                        f_srt.write(f"{start_str} --> {end_str}\n")

                        # 写入转录文本
                        f_srt.write(segment['text'].strip() + '\n\n') # strip() 移除 Whisper 可能生成的首尾空格

                print(f"SRT file saved to: {srt_filepath}")

                # 如果仍然需要txt文件，可以保留以下代码
                # transcription = result['text']
                # with open(os.path.join(dirpath, filename_no_ext + '.txt'), 'w', encoding='utf-8') as f_txt:
                #     f_txt.write(transcription)

                pbar.update(1)

print("\nAll files transcribed and SRT files generated.")

# ... (Whisper 转录代码，获取 result['segments']) ...

# 假设您已通过 PyAnnote 获取了说话人分离结果，
# 格式可能类似：[{'speaker': 'SPEAKER_00', 'start': 0.5, 'end': 2.1}, ...]
# diarization_result = run_pyannote_diarization(filepath) 

# 合并 Whisper 和 Diarization 结果，生成增强型 SRT
# for i, segment in enumerate(result['segments']):
#     # 查找与当前segment时间重叠的diarization_result，确定说话人
#     speaker_label = get_speaker_for_segment(segment, diarization_result) 
#     formatted_text = f"[{speaker_label}] {segment['text'].strip()}"
#     # ... 写入 SRT 文件 ...