This commit is contained in:
Karl 2024-01-08 13:10:49 +00:00
parent 09e718b597
commit f19f886fa0
4 changed files with 3 additions and 55 deletions

View File

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

View File

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

View File

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

View File

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