Merge pull request #3 from karl0ss/reworked

Reworked

.gitignore

@@ -6,3 +6,4 @@ __pycache__
 venv/
 test/
 .vscode/launch.json
+config.cfg

.vscode/launch.json

@@ -12,19 +12,9 @@
             "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": ""
-            }
         }
     ]
 }

README.md

@@ -15,7 +15,7 @@ Clunky, and slow, but works.
 
 ## Usage
 
-The following command will generate a `subtitled/video.mp4` file contained the input video with overlayed subtitles.
+<!-- The following command will generate a `subtitled/video.mp4` file contained the input video with overlayed subtitles.
 
     faster_auto_subtitle /path/to/video.mp4 -o subtitled/
 
@@ -47,4 +47,4 @@ You can use `sample_interval` parameter to generate subtitles for a portion of t
 
 ## License
 
-This script is open-source and licensed under the MIT License. For more details, check the [LICENSE](LICENSE) file.
+This script is open-source and licensed under the MIT License. For more details, check the [LICENSE](LICENSE) file. -->

auto_subtitle/cli.py

@@ -29,12 +29,6 @@ def main():
                         "int16", "float16", "bfloat16", "float32"],
                         help="Type to use for computation. \
                               See https://opennmt.net/CTranslate2/quantization.html.")
-    parser.add_argument("--output_dir", "-o", type=str,
-                        default=".", help="directory to save the outputs")
-    parser.add_argument("--output_srt", type=str2bool, default=False,
-                        help="whether to output the .srt file along with the video files")
-    parser.add_argument("--srt_only", type=str2bool, default=False,
-                        help="only generate the .srt file and not create overlayed video")
     parser.add_argument("--beam_size", type=int, default=5,
                         help="model parameter, tweak to increase accuracy")
     parser.add_argument("--no_speech_threshold", type=float, default=0.6,

auto_subtitle/main.py

@@ -3,7 +3,7 @@ import warnings
 import tempfile
 import time
 from utils.files import filename, write_srt
-from utils.ffmpeg import get_audio, add_subs_new
+from utils.ffmpeg import get_audio, add_subtitles_to_mp4
 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
@@ -11,15 +11,10 @@ from utils.whisper import WhisperAI
 
 def process(args: dict):
     model_name: str = args.pop("model")
-    output_dir: str = args.pop("output_dir")
-    output_srt: bool = args.pop("output_srt")
-    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)
-
     if model_name.endswith(".en"):
         warnings.warn(
             f"{model_name} is an English-only model, forcing English detection.")
@@ -39,19 +34,14 @@ def process(args: dict):
         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)
+        subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args)
 
-        if srt_only:
-            return
-
-        add_subs_new(subtitles)
+        add_subtitles_to_mp4(subtitles)
         update_show_in_soarr(episode['sonarrSeriesId'])
         time.sleep(5)
         sync_series()
         
 
-
 def get_subtitles(audio_paths: list, output_dir: str,
                   model_args: dict, transcribe_args: dict):
     model = WhisperAI(model_args, transcribe_args)

auto_subtitle/utils/bazarr.py

@@ -1,9 +1,13 @@
 import requests
-import os
-token = os.getenv('bazarr_token')
+import configparser
+config = configparser.RawConfigParser()
+config.read('config.cfg')
+
+token = config._sections['bazarr']['token']
+base_url = config._sections['bazarr']['url']
 
 def get_wanted_episodes():
-    url = "http://192.168.4.23/api/episodes/wanted"
+    url = f"{base_url}/api/episodes/wanted"
 
     payload={}
     headers = {
@@ -17,7 +21,7 @@ def get_wanted_episodes():
 
 
 def get_episode_details(episode_id: str):
-    url = f"http://192.168.4.23/api/episodes?episodeid%5B%5D={episode_id}"
+    url = f"{base_url}/api/episodes?episodeid%5B%5D={episode_id}"
 
     payload={}
     headers = {
@@ -30,7 +34,7 @@ def get_episode_details(episode_id: str):
 
 
 def sync_series():
-    url = f"http://192.168.4.23/api/system/tasks?taskid=update_series"
+    url = f"{base_url}/api/system/tasks?taskid=update_series"
 
     payload={}
     headers = {
@@ -39,4 +43,5 @@ def sync_series():
     }
 
     response = requests.request("POST", url, headers=headers, data=payload)
-    return response.json()['data'][0]
+    if response.status_code == 204:
+        print('Updated Bazarr')

auto_subtitle/utils/ffmpeg.py

@@ -1,7 +1,6 @@
 import os
 import tempfile
 import ffmpeg
-from .mytempfile import MyTempFile
 from .files import filename
 
 
@@ -37,7 +36,7 @@ def get_audio(paths: list, audio_channel_index: int, sample_interval: list):
     return audio_paths
 
 
-def add_subs_new(subtitles: dict):
+def add_subtitles_to_mp4(subtitles: dict):
     
     input_file = list(subtitles.keys())[0]
     subtitle_file = subtitles[input_file]
@@ -51,5 +50,6 @@ def add_subs_new(subtitles: dict):
     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)
+    # remove tempfiles
+    os.remove(subtitle_file)
+    os.remove(subtitle_file.replace(".srt",".wav"))

auto_subtitle/utils/mytempfile.py

@@ -1,35 +0,0 @@
-import tempfile
-import os
-import shutil
-
-class MyTempFile:
-    """
-    A context manager for creating a temporary file in current directory, copying the content from
-    a specified file, and handling cleanup operations upon exiting the context.
-
-    Usage:
-    ```python
-    with MyTempFile(file_path) as temp_file_manager:
-        # Access the temporary file using temp_file_manager.tmp_file
-        # ...
-    # The temporary file is automatically closed and removed upon exiting the context.
-    ```
-
-    Args:
-    - file_path (str): The path to the file whose content will be copied to the temporary file.
-    """
-    def __init__(self, file_path):
-        self.file_path = file_path
-        self.tmp_file = None
-        self.tmp_file_path = None
-
-    def __enter__(self):
-        self.tmp_file = tempfile.NamedTemporaryFile('w', dir='.', delete=False)
-        self.tmp_file_path = os.path.relpath(self.tmp_file.name, '.')
-        shutil.copyfile(self.file_path, self.tmp_file_path)
-        return self
-
-    def __exit__(self, exc_type, exc_value, exc_traceback):
-        self.tmp_file.close()
-        if os.path.isfile(self.tmp_file_path):
-            os.remove(self.tmp_file_path)

auto_subtitle/utils/sonarr.py

@@ -1,10 +1,14 @@
 import requests
 import json
-import os
-token = os.getenv('sonarr_token')
+import configparser
+config = configparser.RawConfigParser()
+config.read('config.cfg')
+
+token = config._sections['sonarr']['token']
+base_url = config._sections['sonarr']['url']
 
 def update_show_in_soarr(show_id):
-    url = "http://192.168.4.9:8989/api/v3/command"
+    url = f"{base_url}/api/v3/command"
 
     payload = json.dumps({
     "name": "RefreshSeries",

config.cfg.example

@@ -0,0 +1,6 @@
+[bazarr]  
+url = http://1.1.1.1
+token = djfkjadncdfjkanvfjkvandfj  
+[sonarr]
+url = http://2.2.2.2:8989
+token = dfifdmnajcdnjcvaldnjlk