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

3
.gitignore vendored
View File

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

15
.vscode/launch.json vendored
View File

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

View File

@ -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,

View File

@ -1,9 +1,11 @@
import os
import warnings
import tempfile
import time
from utils.files import filename, write_srt
from utils.ffmpeg import get_audio, overlay_subtitles, add_subs_new
from utils.bazarr import get_wanted_episodes, get_episode_details
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
@ -14,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,27 +27,29 @@ def process(args: dict):
# if translate task used and language argument is set, then use it
elif language != "auto":
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_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
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,

View File

@ -1,6 +1,6 @@
import requests
token = ''
import os
token = os.getenv('bazarr_token')
def get_wanted_episodes():
url = "http://192.168.4.23/api/episodes/wanted"
@ -26,4 +26,17 @@ def get_episode_details(episode_id: str):
}
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]

View File

@ -37,47 +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")
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
def add_subs_new(subtitles: dict):
input_file = list(subtitles.keys())[0]
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)
subtitle_stream = fmp.input(subtitle_file)
input_stream = ffmpeg.input(input_file+'_edit')
subtitle_stream = ffmpeg.input(subtitle_file)
# 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'})
fmp.run(output)
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)

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")