mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 02:48:38 -04:00
127 lines
3.9 KiB
Python
Executable file
127 lines
3.9 KiB
Python
Executable file
#! /usr/bin/env python3
|
|
"""
|
|
Transcribes video files, then creates an vtt file next to each.
|
|
|
|
# ls
|
|
video1.mp4
|
|
video2.mp4
|
|
|
|
# generate_subtitles.py -vv --language en --model small.en video1.mp4 video2.mp4
|
|
|
|
# ls
|
|
video1.mp4
|
|
video1.en.vtt
|
|
video2.mp4
|
|
video2.en.vtt
|
|
"""
|
|
import argparse
|
|
import logging
|
|
import pathlib
|
|
import sys
|
|
from typing import List
|
|
|
|
try:
|
|
# This is not in the requirements.txt because it is not required and does not work on a Raspberry Pi.
|
|
import whisper
|
|
except ImportError:
|
|
whisper = None
|
|
|
|
__all__ = ['video_speech_to_vtt_file', 'video_speech_to_vtt']
|
|
|
|
logger = logging.getLogger(__name__)
|
|
handler = logging.StreamHandler()
|
|
handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
|
|
logger.addHandler(handler)
|
|
|
|
MINUTE = 60
|
|
HOUR = MINUTE * 60
|
|
|
|
|
|
def format_timestamp(seconds: float):
|
|
"""Convert seconds to the expected vtt timestamp format."""
|
|
seconds, milliseconds = divmod(seconds, 1)
|
|
hours, seconds = divmod(seconds, HOUR)
|
|
minutes, seconds = divmod(seconds, MINUTE)
|
|
|
|
milliseconds *= 1000
|
|
hours, minutes, seconds, milliseconds = int(hours), int(minutes), int(seconds), int(milliseconds)
|
|
|
|
timestamp = f'{hours:02d}:{minutes:02d}:{seconds:02d}.{milliseconds:03d}'
|
|
return timestamp
|
|
|
|
|
|
def video_speech_to_vtt(file: pathlib.Path, model, language):
|
|
"""Transcribe a video file. Return the text of the vtt file."""
|
|
logger.info(f'Transcribing: {file}')
|
|
transcription = model.transcribe(str(file), language=language)
|
|
segments = transcription['segments']
|
|
|
|
vtt_lines = []
|
|
for segment in segments:
|
|
start, end, id_, text = segment['start'], segment['end'], segment['id'], segment['text']
|
|
logger.debug(f'{file} === {start} --> {end} {text}')
|
|
start, end = format_timestamp(start), format_timestamp(end)
|
|
vtt_lines.append(f'{id_ + 1}\n{start} --> {end}\n{text}\n')
|
|
|
|
contents = '\n'.join(vtt_lines)
|
|
contents = f'WEBVTT Kind: captions; Language: {language}\n' \
|
|
f'{contents}'
|
|
return contents
|
|
|
|
|
|
def video_speech_to_vtt_file(video_path: pathlib.Path, vtt_path: pathlib.Path, model, language: str):
|
|
"""Transcribe a video, save the resulting transcription in the vtt file."""
|
|
vtt = video_speech_to_vtt(video_path, model, language)
|
|
vtt_path.write_text(vtt, encoding='utf-8')
|
|
|
|
|
|
def main(video_files: List[str], model: str, language: str):
|
|
video_files = [pathlib.Path(i) for i in video_files]
|
|
|
|
# Check that all files are real before we do anything.
|
|
for video_file in video_files:
|
|
if not video_file.is_file():
|
|
print(f'{video_file} is not a video file', file=sys.stderr)
|
|
sys.exit(2)
|
|
|
|
vtt_suffix = f'.{language}.vtt' if language else '.vtt'
|
|
|
|
logger.info(f'Loading model {model}...')
|
|
model = whisper.load_model(model)
|
|
logger.debug('Model loading complete.')
|
|
|
|
for video_file in video_files:
|
|
vtt_path = video_file.with_suffix(vtt_suffix)
|
|
if vtt_path.exists():
|
|
logger.warning(f'{vtt_path} already exists, skipping...')
|
|
continue
|
|
|
|
video_speech_to_vtt_file(video_file, vtt_path, model, args.language)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser('Transcribes video files, then creates an vtt file next to each.')
|
|
parser.add_argument('files', nargs='+')
|
|
parser.add_argument('--model', default='small')
|
|
parser.add_argument('--language', default='en')
|
|
parser.add_argument('-v', action='count', default=0)
|
|
args = parser.parse_args()
|
|
|
|
if whisper is None:
|
|
print('You must install whisper. See https://github.com/openai/whisper', file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
if args.v == 0:
|
|
logger.setLevel(logging.WARNING)
|
|
elif args.v == 1:
|
|
logger.setLevel(logging.INFO)
|
|
elif args.v >= 2:
|
|
logger.setLevel(logging.DEBUG)
|
|
logger.debug(f'Debug logging')
|
|
|
|
main(
|
|
args.files,
|
|
args.model,
|
|
args.language,
|
|
)
|
|
sys.exit(0)
|