commit main

README.md

@@ -1 +1,10 @@
 # auto_ai_subtitle
+
+基于openai/whisper、translate、ffmpeg，自动为视频生成翻译过后的srt字幕文件
+
+## 使用方法
+安装依赖 `pip install -r requirements.txt`
+
+将配置信息填入 *config.ini*
+
+然后执行 `python main.py`

config.yaml

@@ -0,0 +1,6 @@
+input: D:\download\ChainsawMan-03.mp4
+output: D:\download\ChainsawMan-03.mp3
+srt_path: D:\download\ChainsawMan-03.srt
+srt_translate_path: D:\download\ChainsawMan-03-zh.srt
+from: ja
+to: zh

main.py

@@ -0,0 +1,21 @@
+import yaml
+
+from script import translate_tool, audio_tool, whisper_tool
+
+if __name__ == '__main__':
+    with open('config.yaml') as f:
+        config = yaml.load(f.read(), Loader=yaml.FullLoader)
+
+    print("audio extract begin")
+    audio_tool.audio_extract(config['input'], config['output'])
+    print("audio extract success")
+
+    print("whisper begin")
+    whisper_tool.do_whisper(config['output'], config['srt_path'])
+    print("whisper success")
+
+    print("translate begin")
+    translate_tool.do_translate(config['srt_path'], config['srt_translate_path'], config['from'], config['to'])
+    print("translate success")
+
+    print("success")

requirements.txt

@@ -0,0 +1,4 @@
+translate>=3.6.1
+pyyaml>=6.0
+openai-whisper
+ffmpeg-python>=0.2.0

script/audio_tool.py

@@ -0,0 +1,5 @@
+import ffmpeg
+
+
+def audio_extract(input, output):
+    ffmpeg.input(input, vn=None).output(output).run()

script/translate_tool.py

@@ -0,0 +1,25 @@
+from translate import Translator
+import re
+
+
+def __translate(translator, text, n):
+    print("\rtranslate line in: ", n + 1, end="")
+    if not re.match(r"[^->\d\n]", text):
+        return text
+    return translator.translate(text)
+
+
+def translate_file(translator_fun, file1, file2, translator=None):
+    with open(file1, 'r', encoding='utf-8') as f1, open(file2, 'w', encoding='utf-8') as f2:
+        lines = f1.readlines()
+        print("translate file total lines: ", len(lines))
+        f2.writelines([translator_fun(translator, line, n) for n, line in enumerate(lines)])
+
+
+def do_translate(file1, file2, form, to):
+    translator = Translator(from_lang=form, to_lang=to)
+    translate_file(__translate, file1, file2, translator)
+
+
+if __name__ == '__main__':
+    do_translate('test.srt', 'test1.srt', 'ja', 'zh')

script/whisper_tool.py

@@ -0,0 +1,34 @@
+import whisper
+
+
+def reformat_time(second):
+    m, s = divmod(second, 60)
+    h, m = divmod(m, 60)
+    hms = "%02d:%02d:%s" % (h, m, str('%.3f' % s).zfill(6))
+    hms = hms.replace('.', ',')
+    return hms
+
+
+def write_srt(seg, srt_path):
+    with open(srt_path, 'w', encoding='utf-8') as f:
+        write_content = [str(n + 1) + '\n'
+                         + reformat_time(i['start'])
+                         + ' --> '
+                         + reformat_time(i['end']) + '\n'
+                         + i['text'] + '\n\n'
+                         for n, i in enumerate(seg)]
+        f.writelines(write_content)
+
+
+def do_whisper(audio, srt_path):
+    model = whisper.load_model("base")
+    print("whisper working...")
+    result = model.transcribe(audio)
+    print("whisper execute success")
+    print("writing srt file...")
+    write_srt(result['segments'], srt_path)
+    print("write srt success")
+
+
+if __name__ == '__main__':
+    do_whisper("test.mp3", "test.srt")