Skip to content

Commit

Permalink
Add option to process folder with videos
Browse files Browse the repository at this point in the history
  • Loading branch information
Sirozha1337 committed Jun 29, 2024
1 parent 5402c58 commit 11942bb
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 15 deletions.
4 changes: 4 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,10 @@ The following command will generate a `subtitled/video.mp4` file contained the i

faster_auto_subtitle /path/to/video.mp4 -o subtitled/

You can also specify a folder with multiple videos, and it will process all of them:

faster_auto_subtitle /path/to/videos/ -o subtitled/

The default setting (which selects the `small` model) works well for transcribing English. You can optionally use a
bigger model for better results (especially with other languages). The available models
are `tiny`, `tiny.en`, `base`, `base.en`, `small`, `small.en`, `medium`, `medium.en`, `large`, `large-v1`, `large-v2`, `large-v3`.
Expand Down
53 changes: 40 additions & 13 deletions auto_subtitle/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from typing import Optional
from .models.Subtitles import Subtitles, SegmentsIterable
from .utils.files import filename, write_srt
from .utils.ffmpeg import get_audio, add_subtitles, preprocess_audio
from .utils.ffmpeg import get_audio, add_subtitles, preprocess_audio, file_has_audio
from .utils.whisper import WhisperAI
from .translation.easynmt_utils import EasyNMTWrapper

Expand Down Expand Up @@ -39,7 +39,7 @@ def process(args: dict):
"subtitle_type": args.pop("subtitle_type")
}

videos = args.pop('video')
paths_to_process = args.pop('video')
audio_channel = args.pop('audio_channel')
model_args = {
"model_size_or_path": model_name,
Expand All @@ -51,16 +51,42 @@ def process(args: dict):
device=model_args['device']) if target_language != 'en' else None

os.makedirs(output_args["output_dir"], exist_ok=True)
for video in videos:
if video.endswith('.wav'):
audio = preprocess_audio(video, audio_channel, sample_interval)
else:
audio = get_audio(video, audio_channel, sample_interval)
for path_to_process in paths_to_process:
process_path(audio_channel, language, output_args, path_to_process, sample_interval,
target_language, transcribe_model, translate_model)

transcribed, translated = perform_task(video, audio, language, target_language,
transcribe_model, translate_model)

save_result(video, transcribed, translated, sample_interval, output_args)
def process_path(audio_channel, language, output_args, path_to_process, sample_interval,
target_language, transcribe_model, translate_model):
if not os.path.exists(path_to_process):
logger.error("File %s does not exist.", path_to_process)
return

if not os.path.isdir(path_to_process):
process_file(audio_channel, language, output_args, sample_interval, target_language,
transcribe_model, translate_model, path_to_process)
return

logger.info("Processing all files in directory %s", path_to_process)
for file_name in os.listdir(path_to_process):
process_file(audio_channel, language, output_args, sample_interval, target_language,
transcribe_model, translate_model, os.path.join(path_to_process, file_name))


def process_file(audio_channel, language, output_args, sample_interval, target_language,
transcribe_model, translate_model, file_name):
if not file_has_audio(file_name):
logger.info("File %s has no audio, skipping.", file_name)
return

if file_name.endswith('.wav'):
audio = preprocess_audio(file_name, audio_channel, sample_interval)
else:
audio = get_audio(file_name, audio_channel, sample_interval)

transcribed, translated = perform_task(file_name, audio, language, target_language,
transcribe_model, translate_model)
save_result(file_name, transcribed, translated, sample_interval, output_args)


def save_result(video: str, transcribed: Subtitles, translated: Subtitles, sample_interval: list,
Expand All @@ -84,9 +110,7 @@ def perform_task(video: str, audio: str, language: str, target_language: str,
transcribed = get_subtitles(video, audio, transcribe_model)
translated = None

logger.info('Subtitles generated.')
if target_language != 'en':
logger.info('Translating subtitles... This might take a while.')
translated = translate_subtitles(
transcribed, language, target_language, translate_model)

Expand All @@ -99,8 +123,11 @@ def translate_subtitles(subtitles: Subtitles, source_lang: str, target_lang: str
if src_lang == '' or src_lang is None:
src_lang = subtitles.language

segments = list(subtitles.segments)
logger.info('Subtitles generated.')
logger.info('Translating subtitles... This might take a while.')
translated_segments = model.translate(
list(subtitles.segments), src_lang, target_lang)
segments, src_lang, target_lang)

return Subtitles(SegmentsIterable(translated_segments), target_lang)

Expand Down
14 changes: 12 additions & 2 deletions auto_subtitle/utils/ffmpeg.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import os
import tempfile
import logging
from typing import Optional
import ffmpeg
from typing import Optional
from .tempfile import SubtitlesTempFile
from .files import filename
from ..models.Subtitles import Subtitles
Expand Down Expand Up @@ -39,6 +39,16 @@ def get_audio(path: str, audio_channel_index: int, sample_interval: Optional[lis
return output_path


def file_has_audio(path: str) -> bool:
try:
audio_info = ffmpeg.probe(path, select_streams='a')
return 'streams' in audio_info \
and audio_info['streams'] is not None \
and len(audio_info['streams']) > 0
except ffmpeg.Error:
return False


def preprocess_audio(path: str, audio_channel_index: int, sample_interval: Optional[list]) -> str:
if sample_interval is not None or audio_channel_index != 0:
return get_audio(path, audio_channel_index, sample_interval)
Expand Down Expand Up @@ -71,7 +81,7 @@ def add_subtitles(path: str, transcribed: Subtitles, translated: Optional[Subtit
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
# 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 SubtitlesTempFile(transcribed) as transcribed_tmp, SubtitlesTempFile(
Expand Down

0 comments on commit 11942bb

Please sign in to comment.