This commit is contained in:
Karl 2024-01-08 09:20:36 +00:00
parent d1d001114f
commit 4cc5c3e39e
7 changed files with 78 additions and 66 deletions

1
.gitignore vendored
View File

@ -5,3 +5,4 @@ build
__pycache__ __pycache__
venv/ venv/
test/ test/
.vscode/launch.json

15
.vscode/launch.json vendored
View File

@ -12,16 +12,19 @@
"console": "integratedTerminal", "console": "integratedTerminal",
"justMyCode": false, "justMyCode": false,
"args": [ "args": [
"Class of '92 - Out of Their League - S08E03 - Episode 3 HDTV-1080p.mkv", // "Class of '92 - Out of Their League - S08E03 - Episode 3 HDTV-1080p.mkv",
// "--model", "--model",
// "base", "base",
// "--srt_only", // "--srt_only",
// "TRUE", // "TRUE",
"--output_srt", // "--output_srt",
"TRUE", // "TRUE",
"-o", "-o",
"./test" "./test"
] ],
"env": {
"token": ""
}
} }
] ]
} }

View File

@ -14,8 +14,6 @@ def main():
""" """
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter) 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", parser.add_argument("--audio_channel", default="0",
type=int, help="audio channel index to use") type=int, help="audio channel index to use")
parser.add_argument("--sample_interval", type=str2timeinterval, default=None, parser.add_argument("--sample_interval", type=str2timeinterval, default=None,

View File

@ -1,9 +1,11 @@
import os import os
import warnings import warnings
import tempfile import tempfile
import time
from utils.files import filename, write_srt from utils.files import filename, write_srt
from utils.ffmpeg import get_audio, overlay_subtitles, add_subs_new from utils.ffmpeg import get_audio, add_subs_new
from utils.bazarr import get_wanted_episodes, get_episode_details 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 from utils.whisper import WhisperAI
@ -14,6 +16,7 @@ def process(args: dict):
srt_only: bool = args.pop("srt_only") srt_only: bool = args.pop("srt_only")
language: str = args.pop("language") language: str = args.pop("language")
sample_interval: str = args.pop("sample_interval") sample_interval: str = args.pop("sample_interval")
audio_channel: str = args.pop('audio_channel')
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
@ -25,26 +28,28 @@ def process(args: dict):
elif language != "auto": elif language != "auto":
args["language"] = language args["language"] = language
a = get_wanted_episodes()
print(f"Found {a['total']} episodes needing subtitles.")
for episode in a['data']:
episode_data = get_episode_details(episode['sonarrEpisodeId'])
print(episode_data)
audios = get_audio(args.pop("video"), args.pop(
'audio_channel'), sample_interval)
model_args = {} 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["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()
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() srt_output_dir = output_dir if output_srt or srt_only else tempfile.gettempdir()
subtitles = get_subtitles(audios, srt_output_dir, model_args, args) subtitles = get_subtitles(audios, srt_output_dir, model_args, args)
if srt_only: if srt_only:
return return
add_subs_new(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, def get_subtitles(audio_paths: list, output_dir: str,

View File

@ -1,6 +1,6 @@
import requests import requests
import os
token = '' token = os.getenv('bazarr_token')
def get_wanted_episodes(): def get_wanted_episodes():
url = "http://192.168.4.23/api/episodes/wanted" url = "http://192.168.4.23/api/episodes/wanted"
@ -27,3 +27,16 @@ def get_episode_details(episode_id: str):
response = requests.request("GET", url, headers=headers, data=payload) response = requests.request("GET", url, headers=headers, data=payload)
return response.json()['data'][0] 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]

View File

@ -37,47 +37,19 @@ def get_audio(paths: list, audio_channel_index: int, sample_interval: list):
return audio_paths return audio_paths
def overlay_subtitles(subtitles: dict, output_dir: str, sample_interval: list): def add_subs_new(subtitles: dict):
for path, srt_path in subtitles.items():
out_path = os.path.join(output_dir, f"{filename(path)}.mp4")
print(f"Adding subtitles to {filename(path)}...")
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)}.")
def add_subs_new(subtitles: dict, output_dir: str, sample_interval: list):
import ffmpeg as fmp
input_file = list(subtitles.keys())[0] input_file = list(subtitles.keys())[0]
subtitle_file = subtitles[input_file] subtitle_file = subtitles[input_file]
output_file = 'class.mp4' output_file = input_file
os.rename(input_file, input_file+'_edit')
input_stream = fmp.input(input_file) input_stream = ffmpeg.input(input_file+'_edit')
subtitle_stream = fmp.input(subtitle_file) subtitle_stream = ffmpeg.input(subtitle_file)
# Combine input video and subtitle # Combine input video and subtitle
output = fmp.output(input_stream, subtitle_stream, output_dir + '/' + output_file.replace('.mkv','.mp4'), c='copy', **{'c:s': 'mov_text'}, **{'metadata:s:s:0': 'language=eng'}) 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)
fmp.run(output) os.remove(input_file+'_edit')
if '.mkv' in output_file:
os.remove(output_file)

View File

@ -0,0 +1,20 @@
import requests
import json
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': 'f6ea49a75a44469daec03969bdf6764d',
}
response = requests.request("POST", url, headers=headers, data=payload)
if response.status_code != 404:
print("Updated show in Sonarr")