Karl преди 1 година
родител
ревизия
4cc5c3e39e
променени са 7 файла, в които са добавени 78 реда и са изтрити 66 реда
  1. 2 1
      .gitignore
  2. 9 6
      .vscode/launch.json
  3. 0 2
      auto_subtitle/cli.py
  4. 21 16
      auto_subtitle/main.py
  5. 15 2
      auto_subtitle/utils/bazarr.py
  6. 11 39
      auto_subtitle/utils/ffmpeg.py
  7. 20 0
      auto_subtitle/utils/sonarr.py

+ 2 - 1
.gitignore

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

+ 9 - 6
.vscode/launch.json

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

+ 0 - 2
auto_subtitle/cli.py

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

+ 21 - 16
auto_subtitle/main.py

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

+ 15 - 2
auto_subtitle/utils/bazarr.py

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

+ 11 - 39
auto_subtitle/utils/ffmpeg.py

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

+ 20 - 0
auto_subtitle/utils/sonarr.py

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