Alex Yancey
2 years ago
committed by
GitHub
parent
1427a8ab73
commit
8b1db07e9f
@ -0,0 +1,286 @@
|
||||
from __future__ import absolute_import
|
||||
import logging
|
||||
|
||||
from requests import Session
|
||||
|
||||
from subliminal_patch.subtitle import Subtitle
|
||||
from subliminal_patch.providers import Provider
|
||||
from subliminal import __short_version__
|
||||
from subliminal.exceptions import ConfigurationError
|
||||
from subzero.language import Language
|
||||
from subliminal.video import Episode, Movie
|
||||
|
||||
from babelfish.exceptions import LanguageReverseError
|
||||
|
||||
import ffmpeg
|
||||
import functools
|
||||
|
||||
# These are all the languages Whisper supports.
|
||||
# from whisper.tokenizer import LANGUAGES
|
||||
|
||||
whisper_languages = {
|
||||
"en": "english",
|
||||
"zh": "chinese",
|
||||
"de": "german",
|
||||
"es": "spanish",
|
||||
"ru": "russian",
|
||||
"ko": "korean",
|
||||
"fr": "french",
|
||||
"ja": "japanese",
|
||||
"pt": "portuguese",
|
||||
"tr": "turkish",
|
||||
"pl": "polish",
|
||||
"ca": "catalan",
|
||||
"nl": "dutch",
|
||||
"ar": "arabic",
|
||||
"sv": "swedish",
|
||||
"it": "italian",
|
||||
"id": "indonesian",
|
||||
"hi": "hindi",
|
||||
"fi": "finnish",
|
||||
"vi": "vietnamese",
|
||||
"he": "hebrew",
|
||||
"uk": "ukrainian",
|
||||
"el": "greek",
|
||||
"ms": "malay",
|
||||
"cs": "czech",
|
||||
"ro": "romanian",
|
||||
"da": "danish",
|
||||
"hu": "hungarian",
|
||||
"ta": "tamil",
|
||||
"no": "norwegian",
|
||||
"th": "thai",
|
||||
"ur": "urdu",
|
||||
"hr": "croatian",
|
||||
"bg": "bulgarian",
|
||||
"lt": "lithuanian",
|
||||
"la": "latin",
|
||||
"mi": "maori",
|
||||
"ml": "malayalam",
|
||||
"cy": "welsh",
|
||||
"sk": "slovak",
|
||||
"te": "telugu",
|
||||
"fa": "persian",
|
||||
"lv": "latvian",
|
||||
"bn": "bengali",
|
||||
"sr": "serbian",
|
||||
"az": "azerbaijani",
|
||||
"sl": "slovenian",
|
||||
"kn": "kannada",
|
||||
"et": "estonian",
|
||||
"mk": "macedonian",
|
||||
"br": "breton",
|
||||
"eu": "basque",
|
||||
"is": "icelandic",
|
||||
"hy": "armenian",
|
||||
"ne": "nepali",
|
||||
"mn": "mongolian",
|
||||
"bs": "bosnian",
|
||||
"kk": "kazakh",
|
||||
"sq": "albanian",
|
||||
"sw": "swahili",
|
||||
"gl": "galician",
|
||||
"mr": "marathi",
|
||||
"pa": "punjabi",
|
||||
"si": "sinhala",
|
||||
"km": "khmer",
|
||||
"sn": "shona",
|
||||
"yo": "yoruba",
|
||||
"so": "somali",
|
||||
"af": "afrikaans",
|
||||
"oc": "occitan",
|
||||
"ka": "georgian",
|
||||
"be": "belarusian",
|
||||
"tg": "tajik",
|
||||
"sd": "sindhi",
|
||||
"gu": "gujarati",
|
||||
"am": "amharic",
|
||||
"yi": "yiddish",
|
||||
"lo": "lao",
|
||||
"uz": "uzbek",
|
||||
"fo": "faroese",
|
||||
"ht": "haitian creole",
|
||||
"ps": "pashto",
|
||||
"tk": "turkmen",
|
||||
"nn": "nynorsk",
|
||||
"mt": "maltese",
|
||||
"sa": "sanskrit",
|
||||
"lb": "luxembourgish",
|
||||
"my": "myanmar",
|
||||
"bo": "tibetan",
|
||||
"tl": "tagalog",
|
||||
"mg": "malagasy",
|
||||
"as": "assamese",
|
||||
"tt": "tatar",
|
||||
"haw": "hawaiian",
|
||||
"ln": "lingala",
|
||||
"ha": "hausa",
|
||||
"ba": "bashkir",
|
||||
"jw": "javanese",
|
||||
"su": "sundanese",
|
||||
}
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@functools.lru_cache(2)
|
||||
def encode_audio_stream(path, audio_stream_language=None):
|
||||
logger.debug("Encoding audio stream to WAV with ffmpeg")
|
||||
|
||||
try:
|
||||
# This launches a subprocess to decode audio while down-mixing and resampling as necessary.
|
||||
inp = ffmpeg.input(path, threads=0)
|
||||
if audio_stream_language:
|
||||
logger.debug(f"Whisper will only use the {audio_stream_language} audio stream for {path}")
|
||||
inp = inp[f'a:m:language:{audio_stream_language}']
|
||||
|
||||
out, _ = inp.output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=16000) \
|
||||
.run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
|
||||
|
||||
except ffmpeg.Error as e:
|
||||
raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
|
||||
|
||||
logger.debug(f"Finished encoding audio stream in {path} with no errors")
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def whisper_get_language(code, name):
|
||||
# Whisper uses an inconsistent mix of alpha2 and alpha3 language codes
|
||||
try:
|
||||
return Language.fromalpha2(code)
|
||||
except LanguageReverseError:
|
||||
return Language.fromname(name)
|
||||
|
||||
|
||||
def whisper_get_language_reverse(alpha3):
|
||||
# Returns the whisper language code given an alpha3b language
|
||||
for wl in whisper_languages:
|
||||
lan = whisper_get_language(wl, whisper_languages[wl])
|
||||
if lan.alpha3 == alpha3:
|
||||
return wl
|
||||
raise ValueError
|
||||
|
||||
|
||||
class WhisperAISubtitle(Subtitle):
|
||||
'''Whisper AI Subtitle.'''
|
||||
provider_name = 'whisperai'
|
||||
hash_verifiable = False
|
||||
|
||||
def __init__(self, language, video):
|
||||
super(WhisperAISubtitle, self).__init__(language)
|
||||
|
||||
self.video = video
|
||||
self.task = None
|
||||
self.audio_language = None
|
||||
self.force_audio_stream = None
|
||||
|
||||
@property
|
||||
def id(self):
|
||||
return self.video.original_name
|
||||
|
||||
def get_matches(self, video):
|
||||
matches = set()
|
||||
|
||||
if isinstance(video, Episode):
|
||||
matches.update(["series", "season", "episode"])
|
||||
elif isinstance(video, Movie):
|
||||
matches.update(["title"])
|
||||
|
||||
return matches
|
||||
|
||||
|
||||
class WhisperAIProvider(Provider):
|
||||
'''Whisper AI Provider.'''
|
||||
|
||||
languages = set()
|
||||
|
||||
for lan in whisper_languages:
|
||||
languages.update({whisper_get_language(lan, whisper_languages[lan])})
|
||||
|
||||
languages.update(set(Language.rebuild(lang, hi=True) for lang in languages))
|
||||
languages.update(set(Language.rebuild(lang, forced=True) for lang in languages))
|
||||
|
||||
video_types = (Episode, Movie)
|
||||
|
||||
def __init__(self, endpoint=None, timeout=None):
|
||||
if not endpoint:
|
||||
raise ConfigurationError('Whisper Web Service Endpoint must be provided')
|
||||
|
||||
if not timeout:
|
||||
raise ConfigurationError('Whisper Web Service Timeout must be provided')
|
||||
|
||||
self.endpoint = endpoint
|
||||
self.timeout = int(timeout)
|
||||
self.session = None
|
||||
|
||||
def initialize(self):
|
||||
self.session = Session()
|
||||
self.session.headers['User-Agent'] = 'Subliminal/%s' % __short_version__
|
||||
|
||||
def terminate(self):
|
||||
self.session.close()
|
||||
|
||||
|
||||
@functools.lru_cache(2048)
|
||||
def detect_language(self, path) -> Language:
|
||||
out = encode_audio_stream(path)
|
||||
|
||||
r = self.session.post(f"{self.endpoint}/detect-language",
|
||||
params={'encode': 'false'},
|
||||
files={'audio_file': out},
|
||||
timeout=self.timeout)
|
||||
|
||||
logger.info(f"Whisper detected language of {path} as {r.json()['detected_language']}")
|
||||
|
||||
return whisper_get_language(r.json()["language_code"], r.json()["detected_language"])
|
||||
|
||||
def query(self, language, video):
|
||||
if language not in self.languages:
|
||||
return None
|
||||
|
||||
sub = WhisperAISubtitle(language, video)
|
||||
sub.task = "transcribe"
|
||||
|
||||
if video.audio_languages:
|
||||
if language.alpha3 in video.audio_languages:
|
||||
sub.audio_language = language.alpha3
|
||||
if len(list(video.audio_languages)) > 1:
|
||||
sub.force_audio_stream = language.alpha3
|
||||
else:
|
||||
sub.task = "translate"
|
||||
sub.audio_language = list(video.audio_languages)[0]
|
||||
else:
|
||||
# We must detect the language manually
|
||||
detected_lang = self.detect_language(video.original_path)
|
||||
|
||||
if detected_lang != language:
|
||||
sub.task = "translate"
|
||||
|
||||
sub.audio_language = detected_lang.alpha3
|
||||
|
||||
if sub.task == "translate":
|
||||
if language.alpha3 != "eng":
|
||||
logger.info(f"Translation only possible from {language} to English")
|
||||
return None
|
||||
|
||||
logger.debug(f"Whisper ({video.original_path}): {sub.audio_language} -> {language.alpha3} [TASK: {sub.task}]")
|
||||
|
||||
return sub
|
||||
|
||||
def list_subtitles(self, video, languages):
|
||||
subtitles = [self.query(l, video) for l in languages]
|
||||
return [s for s in subtitles if s is not None]
|
||||
|
||||
def download_subtitle(self, subtitle: WhisperAISubtitle):
|
||||
# Invoke Whisper through the API. This may take a long time depending on the file.
|
||||
# TODO: This loads the entire file into memory, find a good way to stream the file in chunks
|
||||
|
||||
out = encode_audio_stream(subtitle.video.original_path, subtitle.force_audio_stream)
|
||||
|
||||
r = self.session.post(f"{self.endpoint}/asr",
|
||||
params={'task': subtitle.task, 'language': whisper_get_language_reverse(subtitle.audio_language), 'output': 'srt', 'encode': 'false'},
|
||||
files={'audio_file': out},
|
||||
timeout=self.timeout)
|
||||
|
||||
subtitle.content = r.content
|
||||
Loading…
Reference in new issue