mirror of
https://github.com/karl0ss/bazarr-ai-sub-generator.git
synced 2025-04-27 07:09:22 +01:00
cleanup
This commit is contained in:
parent
09e718b597
commit
f19f886fa0
@ -29,12 +29,6 @@ def main():
|
|||||||
"int16", "float16", "bfloat16", "float32"],
|
"int16", "float16", "bfloat16", "float32"],
|
||||||
help="Type to use for computation. \
|
help="Type to use for computation. \
|
||||||
See https://opennmt.net/CTranslate2/quantization.html.")
|
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,
|
parser.add_argument("--beam_size", type=int, default=5,
|
||||||
help="model parameter, tweak to increase accuracy")
|
help="model parameter, tweak to increase accuracy")
|
||||||
parser.add_argument("--no_speech_threshold", type=float, default=0.6,
|
parser.add_argument("--no_speech_threshold", type=float, default=0.6,
|
||||||
|
@ -3,7 +3,7 @@ import warnings
|
|||||||
import tempfile
|
import tempfile
|
||||||
import time
|
import time
|
||||||
from utils.files import filename, write_srt
|
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.bazarr import get_wanted_episodes, get_episode_details, sync_series
|
||||||
from utils.sonarr import update_show_in_soarr
|
from utils.sonarr import update_show_in_soarr
|
||||||
from utils.whisper import WhisperAI
|
from utils.whisper import WhisperAI
|
||||||
@ -11,15 +11,10 @@ from utils.whisper import WhisperAI
|
|||||||
|
|
||||||
def process(args: dict):
|
def process(args: dict):
|
||||||
model_name: str = args.pop("model")
|
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")
|
language: str = args.pop("language")
|
||||||
sample_interval: str = args.pop("sample_interval")
|
sample_interval: str = args.pop("sample_interval")
|
||||||
audio_channel: str = args.pop('audio_channel')
|
audio_channel: str = args.pop('audio_channel')
|
||||||
|
|
||||||
# os.makedirs(output_dir, exist_ok=True)
|
|
||||||
|
|
||||||
if model_name.endswith(".en"):
|
if model_name.endswith(".en"):
|
||||||
warnings.warn(
|
warnings.warn(
|
||||||
f"{model_name} is an English-only model, forcing English detection.")
|
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']}")
|
print(f"Processing {episode['seriesTitle']} - {episode['episode_number']}")
|
||||||
episode_data = get_episode_details(episode['sonarrEpisodeId'])
|
episode_data = get_episode_details(episode['sonarrEpisodeId'])
|
||||||
audios = get_audio([episode_data['path']], audio_channel, sample_interval)
|
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)
|
subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args)
|
||||||
|
|
||||||
# if srt_only:
|
add_subtitles_to_mp4(subtitles)
|
||||||
# return
|
|
||||||
|
|
||||||
add_subs_new(subtitles)
|
|
||||||
update_show_in_soarr(episode['sonarrSeriesId'])
|
update_show_in_soarr(episode['sonarrSeriesId'])
|
||||||
time.sleep(5)
|
time.sleep(5)
|
||||||
sync_series()
|
sync_series()
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def get_subtitles(audio_paths: list, output_dir: str,
|
def get_subtitles(audio_paths: list, output_dir: str,
|
||||||
model_args: dict, transcribe_args: dict):
|
model_args: dict, transcribe_args: dict):
|
||||||
model = WhisperAI(model_args, transcribe_args)
|
model = WhisperAI(model_args, transcribe_args)
|
||||||
|
@ -1,7 +1,6 @@
|
|||||||
import os
|
import os
|
||||||
import tempfile
|
import tempfile
|
||||||
import ffmpeg
|
import ffmpeg
|
||||||
from .mytempfile import MyTempFile
|
|
||||||
from .files import filename
|
from .files import filename
|
||||||
|
|
||||||
|
|
||||||
@ -37,7 +36,7 @@ def get_audio(paths: list, audio_channel_index: int, sample_interval: list):
|
|||||||
return audio_paths
|
return audio_paths
|
||||||
|
|
||||||
|
|
||||||
def add_subs_new(subtitles: dict):
|
def add_subtitles_to_mp4(subtitles: dict):
|
||||||
|
|
||||||
input_file = list(subtitles.keys())[0]
|
input_file = list(subtitles.keys())[0]
|
||||||
subtitle_file = subtitles[input_file]
|
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