mirror of
https://github.com/lrnselfreliance/wrolpi.git
synced 2026-10-26 19:01:39 -04:00
74 lines
2.2 KiB
Python
Executable file
74 lines
2.2 KiB
Python
Executable file
#! /usr/bin/env python3
|
|
from pathlib import Path
|
|
from typing import Generator, List, Union
|
|
|
|
import srt
|
|
import webvtt
|
|
|
|
from wrolpi.common import logger, chunks
|
|
from wrolpi.db import get_db_session
|
|
from wrolpi.errors import UnknownCaptionFile
|
|
from .common import get_absolute_video_caption
|
|
from .models import Video
|
|
|
|
|
|
def get_caption_text(caption_path: Union[str, Path]) -> Generator:
|
|
"""
|
|
Return all text from each caption of a caption file.
|
|
"""
|
|
if str(caption_path).endswith('vtt'):
|
|
for caption in webvtt.read(caption_path):
|
|
text = str(caption.text).strip()
|
|
yield text
|
|
else:
|
|
with open(caption_path, 'rt') as fh:
|
|
contents = fh.read()
|
|
for subtitle in srt.parse(contents):
|
|
yield subtitle.content
|
|
|
|
|
|
def get_unique_caption_lines(caption_path: Union[str, Path]) -> Generator:
|
|
"""
|
|
Return all unique lines from each caption of a caption file.
|
|
"""
|
|
last_line = None
|
|
for text in get_caption_text(caption_path):
|
|
for line in text.split('\n'):
|
|
if line and line != last_line:
|
|
last_line = line
|
|
yield line
|
|
|
|
|
|
def process_captions(video: Video):
|
|
"""
|
|
Parse and insert captions for a video record.
|
|
"""
|
|
caption_path = get_absolute_video_caption(video)
|
|
if not caption_path.exists():
|
|
raise UnknownCaptionFile()
|
|
try:
|
|
lines = get_unique_caption_lines(str(caption_path))
|
|
block = '\n'.join(lines)
|
|
video.caption = block
|
|
return
|
|
except UnicodeDecodeError:
|
|
# Failed to decode the caption file
|
|
# TODO handle this error
|
|
pass
|
|
except webvtt.errors.MalformedFileError:
|
|
# File format is broken
|
|
pass
|
|
except webvtt.errors.MalformedCaptionError:
|
|
# Captions form is broken somehow
|
|
pass
|
|
|
|
logger.debug(f'Failed to parse caption file {caption_path}')
|
|
|
|
|
|
async def insert_bulk_captions(video_ids: List[int]):
|
|
for video_ids in chunks(video_ids, 10):
|
|
with get_db_session(commit=True) as session:
|
|
for idx, video_id in enumerate(video_ids):
|
|
video = session.query(Video).filter_by(id=video_id).one()
|
|
process_captions(video)
|
|
logger.debug(f'Inserted {len(video_ids)} captions')
|