mirror of
https://github.com/karl0ss/bazarr-ai-sub-generator.git
synced 2025-04-26 14:59:21 +01:00
commit
816b132d5d
5
.gitignore
vendored
5
.gitignore
vendored
@ -2,4 +2,7 @@ dist
|
||||
.DS_Store
|
||||
*.egg-info
|
||||
build
|
||||
__pycache__
|
||||
__pycache__
|
||||
venv/
|
||||
test/
|
||||
.vscode/launch.json
|
||||
|
30
.vscode/launch.json
vendored
Normal file
30
.vscode/launch.json
vendored
Normal file
@ -0,0 +1,30 @@
|
||||
{
|
||||
// Use IntelliSense to learn about possible attributes.
|
||||
// Hover to view descriptions of existing attributes.
|
||||
// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
|
||||
"version": "0.2.0",
|
||||
"configurations": [
|
||||
{
|
||||
"name": "Python: Current File",
|
||||
"type": "python",
|
||||
"request": "launch",
|
||||
"program": "${file}",
|
||||
"console": "integratedTerminal",
|
||||
"justMyCode": false,
|
||||
"args": [
|
||||
// "Class of '92 - Out of Their League - S08E03 - Episode 3 HDTV-1080p.mkv",
|
||||
"--model",
|
||||
"base",
|
||||
// "--srt_only",
|
||||
// "TRUE",
|
||||
// "--output_srt",
|
||||
// "TRUE",
|
||||
"-o",
|
||||
"./test"
|
||||
],
|
||||
"env": {
|
||||
"token": ""
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
@ -1,8 +1,8 @@
|
||||
import argparse
|
||||
from faster_whisper import available_models
|
||||
from .utils.constants import LANGUAGE_CODES
|
||||
from .main import process
|
||||
from .utils.convert import str2bool, str2timeinterval
|
||||
from utils.constants import LANGUAGE_CODES
|
||||
from main import process
|
||||
from utils.convert import str2bool, str2timeinterval
|
||||
|
||||
|
||||
def main():
|
||||
@ -14,8 +14,6 @@ def main():
|
||||
"""
|
||||
parser = argparse.ArgumentParser(
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument("video", nargs="+", type=str,
|
||||
help="paths to video files to transcribe")
|
||||
parser.add_argument("--audio_channel", default="0",
|
||||
type=int, help="audio channel index to use")
|
||||
parser.add_argument("--sample_interval", type=str2timeinterval, default=None,
|
||||
|
@ -1,9 +1,12 @@
|
||||
import os
|
||||
import warnings
|
||||
import tempfile
|
||||
from .utils.files import filename, write_srt
|
||||
from .utils.ffmpeg import get_audio, overlay_subtitles
|
||||
from .utils.whisper import WhisperAI
|
||||
import time
|
||||
from utils.files import filename, write_srt
|
||||
from utils.ffmpeg import get_audio, add_subs_new
|
||||
from utils.bazarr import get_wanted_episodes, get_episode_details, sync_series
|
||||
from utils.sonarr import update_show_in_soarr
|
||||
from utils.whisper import WhisperAI
|
||||
|
||||
|
||||
def process(args: dict):
|
||||
@ -13,6 +16,7 @@ def process(args: dict):
|
||||
srt_only: bool = args.pop("srt_only")
|
||||
language: str = args.pop("language")
|
||||
sample_interval: str = args.pop("sample_interval")
|
||||
audio_channel: str = args.pop('audio_channel')
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
@ -24,21 +28,28 @@ def process(args: dict):
|
||||
elif language != "auto":
|
||||
args["language"] = language
|
||||
|
||||
audios = get_audio(args.pop("video"), args.pop(
|
||||
'audio_channel'), sample_interval)
|
||||
|
||||
model_args = {}
|
||||
model_args["model_size_or_path"] = model_name
|
||||
model_args["device"] = args.pop("device")
|
||||
model_args["compute_type"] = args.pop("compute_type")
|
||||
|
||||
list_of_episodes_needing_subtitles = get_wanted_episodes()
|
||||
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)
|
||||
srt_output_dir = output_dir if output_srt or srt_only else tempfile.gettempdir()
|
||||
subtitles = get_subtitles(audios, srt_output_dir, model_args, args)
|
||||
|
||||
srt_output_dir = output_dir if output_srt or srt_only else tempfile.gettempdir()
|
||||
subtitles = get_subtitles(audios, srt_output_dir, model_args, args)
|
||||
if srt_only:
|
||||
return
|
||||
|
||||
if srt_only:
|
||||
return
|
||||
|
||||
overlay_subtitles(subtitles, output_dir, sample_interval)
|
||||
add_subs_new(subtitles)
|
||||
update_show_in_soarr(episode['sonarrSeriesId'])
|
||||
time.sleep(5)
|
||||
sync_series()
|
||||
|
||||
|
||||
|
||||
def get_subtitles(audio_paths: list, output_dir: str,
|
||||
|
42
auto_subtitle/utils/bazarr.py
Normal file
42
auto_subtitle/utils/bazarr.py
Normal file
@ -0,0 +1,42 @@
|
||||
import requests
|
||||
import os
|
||||
token = os.getenv('bazarr_token')
|
||||
|
||||
def get_wanted_episodes():
|
||||
url = "http://192.168.4.23/api/episodes/wanted"
|
||||
|
||||
payload={}
|
||||
headers = {
|
||||
'accept': 'application/json',
|
||||
'X-API-KEY': token
|
||||
}
|
||||
|
||||
response = requests.request("GET", url, headers=headers, data=payload)
|
||||
|
||||
return response.json()
|
||||
|
||||
|
||||
def get_episode_details(episode_id: str):
|
||||
url = f"http://192.168.4.23/api/episodes?episodeid%5B%5D={episode_id}"
|
||||
|
||||
payload={}
|
||||
headers = {
|
||||
'accept': 'application/json',
|
||||
'X-API-KEY': token
|
||||
}
|
||||
|
||||
response = requests.request("GET", url, headers=headers, data=payload)
|
||||
return response.json()['data'][0]
|
||||
|
||||
|
||||
def sync_series():
|
||||
url = f"http://192.168.4.23/api/system/tasks?taskid=update_series"
|
||||
|
||||
payload={}
|
||||
headers = {
|
||||
'accept': 'application/json',
|
||||
'X-API-KEY': token
|
||||
}
|
||||
|
||||
response = requests.request("POST", url, headers=headers, data=payload)
|
||||
return response.json()['data'][0]
|
@ -37,32 +37,19 @@ def get_audio(paths: list, audio_channel_index: int, sample_interval: list):
|
||||
return audio_paths
|
||||
|
||||
|
||||
def overlay_subtitles(subtitles: dict, output_dir: str, sample_interval: list):
|
||||
for path, srt_path in subtitles.items():
|
||||
out_path = os.path.join(output_dir, f"{filename(path)}.mp4")
|
||||
def add_subs_new(subtitles: dict):
|
||||
|
||||
input_file = list(subtitles.keys())[0]
|
||||
subtitle_file = subtitles[input_file]
|
||||
output_file = input_file
|
||||
os.rename(input_file, input_file+'_edit')
|
||||
|
||||
print(f"Adding subtitles to {filename(path)}...")
|
||||
input_stream = ffmpeg.input(input_file+'_edit')
|
||||
subtitle_stream = ffmpeg.input(subtitle_file)
|
||||
|
||||
ffmpeg_input_args = {}
|
||||
if sample_interval is not None:
|
||||
ffmpeg_input_args['ss'] = str(sample_interval[0])
|
||||
|
||||
ffmpeg_output_args = {}
|
||||
if sample_interval is not None:
|
||||
ffmpeg_output_args['t'] = str(
|
||||
sample_interval[1] - sample_interval[0])
|
||||
|
||||
# HACK: On Windows it's impossible to use absolute subtitle file path with ffmpeg
|
||||
# so we use temp copy instead
|
||||
# see: https://github.com/kkroening/ffmpeg-python/issues/745
|
||||
with MyTempFile(srt_path) as srt_temp:
|
||||
video = ffmpeg.input(path, **ffmpeg_input_args)
|
||||
audio = video.audio
|
||||
|
||||
ffmpeg.concat(
|
||||
video.filter(
|
||||
'subtitles', srt_temp.tmp_file_path,
|
||||
force_style="OutlineColour=&H40000000,BorderStyle=3"), audio, v=1, a=1
|
||||
).output(out_path, **ffmpeg_output_args).run(quiet=True, overwrite_output=True)
|
||||
|
||||
print(f"Saved subtitled video to {os.path.abspath(out_path)}.")
|
||||
# Combine input video and subtitle
|
||||
output = ffmpeg.output(input_stream, subtitle_stream, output_file.replace('.mkv','.mp4'), c='copy', **{'c:s': 'mov_text'}, **{'metadata:s:s:0': 'language=eng'})
|
||||
ffmpeg.run(output, quiet=True, overwrite_output=True)
|
||||
os.remove(input_file+'_edit')
|
||||
if '.mkv' in output_file:
|
||||
os.remove(output_file)
|
21
auto_subtitle/utils/sonarr.py
Normal file
21
auto_subtitle/utils/sonarr.py
Normal file
@ -0,0 +1,21 @@
|
||||
import requests
|
||||
import json
|
||||
import os
|
||||
token = os.getenv('sonarr_token')
|
||||
|
||||
def update_show_in_soarr(show_id):
|
||||
url = "http://192.168.4.9:8989/api/v3/command"
|
||||
|
||||
payload = json.dumps({
|
||||
"name": "RefreshSeries",
|
||||
"seriesId": show_id
|
||||
})
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'X-Api-Key': token,
|
||||
}
|
||||
|
||||
response = requests.request("POST", url, headers=headers, data=payload)
|
||||
|
||||
if response.status_code != 404:
|
||||
print("Updated show in Sonarr")
|
Loading…
x
Reference in New Issue
Block a user