mirror of
https://github.com/karl0ss/bazarr-ai-sub-generator.git
synced 2025-04-26 14:59:21 +01:00
cwc
This commit is contained in:
parent
d1d001114f
commit
4cc5c3e39e
1
.gitignore
vendored
1
.gitignore
vendored
@ -5,3 +5,4 @@ build
|
|||||||
__pycache__
|
__pycache__
|
||||||
venv/
|
venv/
|
||||||
test/
|
test/
|
||||||
|
.vscode/launch.json
|
||||||
|
15
.vscode/launch.json
vendored
15
.vscode/launch.json
vendored
@ -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": ""
|
||||||
|
}
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
}
|
}
|
@ -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,
|
||||||
|
@ -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,
|
||||||
|
@ -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]
|
@ -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)
|
20
auto_subtitle/utils/sonarr.py
Normal file
20
auto_subtitle/utils/sonarr.py
Normal 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")
|
Loading…
x
Reference in New Issue
Block a user