forked from bubblesub/bubblesub
-
Notifications
You must be signed in to change notification settings - Fork 0
/
example_plugin.py
103 lines (90 loc) · 3.41 KB
/
example_plugin.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
# Example plugin: speech recognition of selected lines
import argparse
import asyncio
import concurrent.futures
import io
import typing as T
import speech_recognition as sr
from bubblesub.api import Api
from bubblesub.api.cmd import BaseCommand
from bubblesub.cfg.menu import MenuCommand, SubMenu
from bubblesub.cmd.common import SubtitlesSelection
from bubblesub.fmt.ass.event import AssEvent
class SpeechRecognitionCommand(BaseCommand):
names = ["sr", "google-speech-recognition"]
help_text = (
"Puts results of Google speech recognition "
"for selected subtitles into their notes."
)
@property
def is_enabled(self) -> bool:
return (
self.args.target.makes_sense
and self.api.audio.current_stream
and self.api.audio.current_stream.is_ready
)
async def run(self) -> None:
await asyncio.get_event_loop().run_in_executor(
None,
self.run_in_background,
await self.args.target.get_subtitles(),
)
def run_in_background(self, subs: T.List[AssEvent]) -> None:
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
future_to_sub = {
executor.submit(self.recognize, sub): sub for sub in subs
}
completed, non_completed = concurrent.futures.wait(
future_to_sub, timeout=5
)
with self.api.undo.capture():
for future, sub in future_to_sub.items():
if future not in completed:
continue
try:
note = future.result()
except sr.UnknownValueError:
self.api.log.warn(f"line #{sub.number}: not recognized")
except sr.RequestError as ex:
self.api.log.error(f"line #{sub.number}: error ({ex})")
else:
self.api.log.info(f"line #{sub.number}: OK")
if sub.note:
sub.note += r"\N" + note
else:
sub.note = note
for future, sub in future_to_sub.items():
if future in non_completed:
self.api.log.info(f"line #{sub.number}: timeout")
def recognize(self, sub: AssEvent) -> str:
self.api.log.info(f"line #{sub.number} - analyzing")
recognizer = sr.Recognizer()
with io.BytesIO() as handle:
self.api.audio.current_stream.save_wav(handle, sub.start, sub.end)
handle.seek(0, io.SEEK_SET)
with sr.AudioFile(handle) as source:
audio = recognizer.record(source)
return recognizer.recognize_google(audio, language=self.args.code)
@staticmethod
def decorate_parser(api: Api, parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"-t",
"--target",
help="subtitles to process",
type=lambda value: SubtitlesSelection(api, value),
default="selected",
)
parser.add_argument("code", help="language code")
COMMANDS = [SpeechRecognitionCommand]
MENU = [
SubMenu(
"&Speech recognition",
[
MenuCommand("&Japanese", "sr ja"),
MenuCommand("&German", "sr de"),
MenuCommand("&French", "sr fr"),
MenuCommand("&Italian", "sr it"),
MenuCommand("&Auto", "sr auto"),
],
)
]