mirror of
https://github.com/karl0ss/bazarr-ai-sub-generator.git
synced 2025-04-26 14:59:21 +01:00
cleanup
This commit is contained in:
parent
09e718b597
commit
f19f886fa0
@ -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,
|
||||
|
@ -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)
|
||||
|
@ -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]
|
||||
|
@ -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)
|
Loading…
x
Reference in New Issue
Block a user