Karl/bazarr-ai-sub-generator
1
0
Fork 0
mirror of https://github.com/karl0ss/bazarr-ai-sub-generator.git synced 2025-04-26 14:59:21 +01:00
Code Issues Packages Projects Releases Wiki Activity

cleanup and add ability to process specifc show only

Browse Source
This commit is contained in:
Karl 2024-07-13 09:35:19 +00:00
parent fde1b4d89e
commit 7e83e4ef1e
4 changed files with 64 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.vscode/launch.json vendored
View File

@ -14,6 +14,8 @@
"args": [
"--model",
"base",
"--show",
"Gary Neville's Soccerbox"
],
}
]

96
bazarr-ai-sub-generator/cli.py
View File

@ -15,16 +15,16 @@ def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument(
"--audio_channel", default="0", type=int, help="audio channel index to use"
)
parser.add_argument(
"--sample_interval",
type=str2timeinterval,
default=None,
help="generate subtitles for a specific \
fragment of the video (e.g. 01:02:05-01:03:45)",
)
# parser.add_argument(
# "--audio_channel", default="0", type=int, help="audio channel index to use"
# )
# parser.add_argument(
# "--sample_interval",
# type=str2timeinterval,
# default=None,
# help="generate subtitles for a specific \
# fragment of the video (e.g. 01:02:05-01:03:45)",
# )
parser.add_argument(
"--model",
default="small",
@ -38,46 +38,46 @@ def main():
choices=["cpu", "cuda", "auto"],
help='Device to use for computation ("cpu", "cuda", "auto")',
)
# parser.add_argument(
# "--compute_type",
# type=str,
# default="default",
# choices=[
# "int8",
# "int8_float32",
# "int8_float16",
# "int8_bfloat16",
# "int16",
# "float16",
# "bfloat16",
# "float32",
# ],
# help="Type to use for computation. \
# See https://opennmt.net/CTranslate2/quantization.html.",
# )
# parser.add_argument(
# "--beam_size",
# type=int,
# default=5,
# help="model parameter, tweak to increase accuracy",
# )
# parser.add_argument(
# "--no_speech_threshold",
# type=float,
# default=0.6,
# help="model parameter, tweak to increase accuracy",
# )
# parser.add_argument(
# "--condition_on_previous_text",
# type=str2bool,
# default=True,
# help="model parameter, tweak to increase accuracy",
# )
parser.add_argument(
"--compute_type",
"--show",
type=str,
default="default",
choices=[
"int8",
"int8_float32",
"int8_float16",
"int8_bfloat16",
"int16",
"float16",
"bfloat16",
"float32",
],
help="Type to use for computation. \
See https://opennmt.net/CTranslate2/quantization.html.",
)
parser.add_argument(
"--beam_size",
type=int,
default=5,
help="model parameter, tweak to increase accuracy",
)
parser.add_argument(
"--no_speech_threshold",
type=float,
default=0.6,
help="model parameter, tweak to increase accuracy",
)
parser.add_argument(
"--condition_on_previous_text",
type=str2bool,
default=True,
help="model parameter, tweak to increase accuracy",
)
parser.add_argument(
"--task",
type=str,
default="transcribe",
choices=["transcribe", "translate"],
default=None,
#choices=["transcribe", "translate"],
help="whether to perform X->X speech recognition ('transcribe') \
or X->English translation ('translate')",
)

13
bazarr-ai-sub-generator/main.py
View File

@ -24,8 +24,9 @@ def process(args: dict):
model_name: str = args.pop("model")
language: str = args.pop("language")
sample_interval: str = args.pop("sample_interval")
audio_channel: str = args.pop("audio_channel")
show: str = args.pop("show")
# sample_interval: str = args.pop("sample_interval")
# audio_channel: str = args.pop("audio_channel")
if model_name.endswith(".en"):
warnings.warn(
@ -37,18 +38,18 @@ def process(args: dict):
args["language"] = language
model_args = {}
model_args["model_size_or_path"] = model_name
# model_args["model_size_or_path"] = model_name
model_args["device"] = args.pop("device")
model_args["compute_type"] = args.pop("compute_type")
# model_args["compute_type"] = args.pop("compute_type")
list_of_episodes_needing_subtitles = get_wanted_episodes()
list_of_episodes_needing_subtitles = get_wanted_episodes(show)
print(
f"Found {list_of_episodes_needing_subtitles['total']} episodes needing subtitles."
)
for episode in list_of_episodes_needing_subtitles["data"]:
print(f"Processing {episode['seriesTitle']} - {episode['episode_number']}")
episode_data = get_episode_details(episode["sonarrEpisodeId"])
audios = get_audio([episode_data["path"]], audio_channel, sample_interval)
audios = get_audio([episode_data["path"]], 0, None)
subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args)
add_subtitles_to_mp4(subtitles)

10
bazarr-ai-sub-generator/utils/bazarr.py
View File

@ -8,15 +8,19 @@ token = config._sections["bazarr"]["token"]
base_url = config._sections["bazarr"]["url"]
def get_wanted_episodes():
def get_wanted_episodes(show: str=None):
url = f"{base_url}/api/episodes/wanted"
payload = {}
headers = {"accept": "application/json", "X-API-KEY": token}
response = requests.request("GET", url, headers=headers, data=payload)
return response.json()
data = response.json()
if show != None:
data['data'] = [item for item in data['data'] if item['seriesTitle'] == show]
data['total'] = len(data['data'])
return data
def get_episode_details(episode_id: str):