Merge pull request #1 from karl0ss/reworked

Reworked

.gitignore

@@ -3,3 +3,6 @@ dist
 *.egg-info
 build
 __pycache__
+venv/
+test/
+.vscode/launch.json

.vscode/launch.json

@@ -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": ""
+            }
+        }
+    ]
+}

auto_subtitle/cli.py

@@ -1,8 +1,8 @@
 import argparse
 from faster_whisper import available_models
-from .utils.constants import LANGUAGE_CODES
+from utils.constants import LANGUAGE_CODES
-from .main import process
+from main import process
-from .utils.convert import str2bool, str2timeinterval
+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,

auto_subtitle/main.py

@@ -1,9 +1,12 @@
 import os
 import warnings
 import tempfile
-from .utils.files import filename, write_srt
+import time
-from .utils.ffmpeg import get_audio, overlay_subtitles
+from utils.files import filename, write_srt
-from .utils.whisper import WhisperAI
+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")
         
-    srt_output_dir = output_dir if output_srt or srt_only else tempfile.gettempdir()
+    list_of_episodes_needing_subtitles = get_wanted_episodes()
-    subtitles = get_subtitles(audios, srt_output_dir, model_args, args)
+    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)
 
-    if srt_only:
+        if srt_only:
-        return
+            return
+
+        add_subs_new(subtitles)
+        update_show_in_soarr(episode['sonarrSeriesId'])
+        time.sleep(5)
+        sync_series()
         
-    overlay_subtitles(subtitles, output_dir, sample_interval)
 
 
 def get_subtitles(audio_paths: list, output_dir: str,

auto_subtitle/utils/bazarr.py

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

auto_subtitle/utils/ffmpeg.py

@@ -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):
+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)}...")
+    input_file = list(subtitles.keys())[0]
+    subtitle_file = subtitles[input_file]
+    output_file = input_file
+    os.rename(input_file, input_file+'_edit')
 
-        ffmpeg_input_args = {}
+    input_stream = ffmpeg.input(input_file+'_edit')
-        if sample_interval is not None:
+    subtitle_stream = ffmpeg.input(subtitle_file)
-            ffmpeg_input_args['ss'] = str(sample_interval[0])
 
-        ffmpeg_output_args = {}
+    # Combine input video and subtitle
-        if sample_interval is not None:
+    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_output_args['t'] = str(
+    ffmpeg.run(output, quiet=True, overwrite_output=True)
-                sample_interval[1] - sample_interval[0])
+    os.remove(input_file+'_edit')
-
+    if '.mkv' in output_file:
-        # HACK: On Windows it's impossible to use absolute subtitle file path with ffmpeg
+        os.remove(output_file)
-        # 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)}.")

auto_subtitle/utils/sonarr.py

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