Move functions related to subtitles OCR in a dedicated module.
This commit is contained in:
Frédéric Tronel
@@ -0,0 +1,186 @@
|
||||
# SPDX-License-Identifier: GPL-2.0-or-later
|
||||
#
|
||||
# Copyright (C) 2026 Frédéric Tronel
|
||||
|
||||
|
||||
import logging
|
||||
from typing import IO
|
||||
from io import TextIOWrapper
|
||||
from datetime import timedelta
|
||||
from os import (
|
||||
read,
|
||||
SEEK_SET,
|
||||
lseek,
|
||||
memfd_create,
|
||||
write,
|
||||
fstat,
|
||||
)
|
||||
import re
|
||||
import os
|
||||
from subprocess import PIPE, Popen
|
||||
|
||||
from iso639 import Lang
|
||||
from iso639.exceptions import InvalidLanguageValue
|
||||
from typeguard import typechecked
|
||||
from tqdm import tqdm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@typechecked
|
||||
def get_tesseract_supported_lang(tesseract_path:str) -> dict[Lang, str]|None:
|
||||
"""
|
||||
Retrieves the set of natural languages supported by the Tesseract OCR tool.
|
||||
|
||||
This function runs the Tesseract binary with the --list-langs option and parses the output
|
||||
to extract the supported languages.
|
||||
|
||||
Args:
|
||||
tesseract_path (str): The path to the Tesseract binary.
|
||||
|
||||
Returns:
|
||||
dict[Lang, str] | None:
|
||||
- A dictionary mapping Lang objects to their corresponding language codes
|
||||
(e.g., "eng" for English)
|
||||
- None if an error occurs while running the Tesseract binary
|
||||
"""
|
||||
res = {}
|
||||
|
||||
with Popen([tesseract_path, '--list-langs'], stdout=PIPE) as tesseract:
|
||||
for line in tesseract.stdout:
|
||||
line = line.decode('utf8')
|
||||
p = re.compile('(?P<lang>[a-z]{3})\n')
|
||||
m = re.match(p,line)
|
||||
if m is not None:
|
||||
try:
|
||||
lang = m.group('lang')
|
||||
key = Lang(lang)
|
||||
res[key] = lang
|
||||
except InvalidLanguageValue as e:
|
||||
logger.warning('Invalid language: %s', e)
|
||||
|
||||
tesseract.wait()
|
||||
|
||||
if tesseract.returncode != 0:
|
||||
logger.error("Tesseract returns an error code: %d",tesseract.returncode)
|
||||
return None
|
||||
|
||||
return res
|
||||
|
||||
|
||||
|
||||
@typechecked
|
||||
def extract_srt(mkvextract:str, filename:str, subtitles:dict[str, list[int]],
|
||||
langs:dict[Lang,str]) -> list[tuple[str,str,str,str]]|None:
|
||||
params = [mkvextract, filename, 'tracks']
|
||||
|
||||
res = []
|
||||
|
||||
for lang in subtitles:
|
||||
iso = Lang(lang)
|
||||
|
||||
if iso in langs:
|
||||
ocrlang = langs[iso]
|
||||
else:
|
||||
logger.warning("Language not supported by Tesseract: %s", iso.name)
|
||||
ocrlang ='osd'
|
||||
|
||||
if len(subtitles[lang]) == 1:
|
||||
params.append(f'{subtitles[lang][0]:d}:{lang}')
|
||||
res.append((f'{lang}.idx', f'{lang}.sub', lang, ocrlang))
|
||||
else:
|
||||
count = 1
|
||||
for track in subtitles[lang]:
|
||||
params.append(f'{track:d}:{lang}-{count:d}')
|
||||
res.append((f'{lang}-{count:d}.idx', f'{lang}-{count:d}.sub', lang, ocrlang))
|
||||
count = count+1
|
||||
|
||||
logger.debug('Executing %s', params)
|
||||
|
||||
env = {**os.environ, 'LANG': 'C'}
|
||||
with Popen(params, stdout=PIPE, close_fds=False, env=env) as extract:
|
||||
pb = tqdm(TextIOWrapper(extract.stdout, encoding="utf-8"), total=100, unit='%',
|
||||
desc='Extraction:')
|
||||
for line in pb:
|
||||
if line.startswith('Progress :'):
|
||||
p = re.compile('^Progress : (?P<progress>[0-9]{1,3})%$')
|
||||
m = p.match(line)
|
||||
if m is None:
|
||||
logger.error('Impossible to parse progress')
|
||||
pb.update(int(m['progress'])-pb.n)
|
||||
pb.update(100-pb.n)
|
||||
pb.refresh()
|
||||
pb.close()
|
||||
|
||||
extract.wait()
|
||||
|
||||
# mkvextract returns 0, 1 or 2 as error code.
|
||||
match extract.returncode:
|
||||
case 0:
|
||||
logger.info('Subtitle tracks were succesfully extracted.')
|
||||
case 1:
|
||||
logger.warning('Mkvextract returns warning')
|
||||
case 2:
|
||||
logger.error('Mkvextract returns an error code: %d', extract.returncode)
|
||||
res = None
|
||||
|
||||
return res
|
||||
|
||||
@typechecked
|
||||
def do_ocr(vobsubocr:str, idxs: list[tuple[str,str,str,str]], duration:timedelta,
|
||||
temporaries:list[IO[bytes]], dump_mem_fd:bool=False):
|
||||
res = []
|
||||
|
||||
for idx_name, _, lang, iso in idxs:
|
||||
srtname = f'{os.path.splitext(idx_name)[0]}.srt'
|
||||
# Tesseract seems to recognize the three dots ... as "su"
|
||||
ldots = re.compile('^su\n$')
|
||||
# Timestamps produced by vobsubocr: 01:52:19,861 --> 01:52:21,641
|
||||
timestamps = re.compile((r'^[0-9]{2}:[0-9]{2}:[0-9]{2},[0-9]{3} \-\-> (?P<hours>[0-9]{2}):'
|
||||
r'(?P<minutes>[0-9]{2}):(?P<seconds>[0-9]{2}),[0-9]{3}$'))
|
||||
|
||||
srtfd = memfd_create(srtname, flags=0)
|
||||
with Popen([vobsubocr, '--lang', iso, idx_name], stdout=PIPE) as ocr:
|
||||
pb = tqdm(TextIOWrapper(ocr.stdout, encoding="utf-8"), total=
|
||||
int(duration/timedelta(seconds=1)), unit='s', desc='OCR')
|
||||
for line in pb:
|
||||
m = re.match(ldots,line)
|
||||
if m is not None:
|
||||
write(srtfd, '...'.encode(encoding='UTF-8'))
|
||||
else:
|
||||
write(srtfd, line.encode(encoding='UTF-8'))
|
||||
|
||||
m = re.match(timestamps, line)
|
||||
if m is not None:
|
||||
hours = int(m.group('hours'))
|
||||
minutes = int(m.group('hours'))
|
||||
seconds = int(m.group('seconds'))
|
||||
ts = timedelta(hours=hours, minutes=minutes, seconds=seconds)
|
||||
pb.n = int(ts/timedelta(seconds=1))
|
||||
pb.update()
|
||||
|
||||
status = ocr.wait()
|
||||
|
||||
if status != 0:
|
||||
logger.error('OCR failed with status code: %d', status)
|
||||
|
||||
if dump_mem_fd:
|
||||
try:
|
||||
with open(srtname,'w', encoding='utf8') as dump_srt:
|
||||
lseek(srtfd, 0, SEEK_SET)
|
||||
srt_length = fstat(srtfd).st_size
|
||||
buf = read(srtfd, srt_length)
|
||||
outfd = dump_srt.fileno()
|
||||
pos = 0
|
||||
while pos < srt_length:
|
||||
pos+=write(outfd, buf[pos:])
|
||||
temporaries.append(dump_srt)
|
||||
except OSError:
|
||||
logger.error('Impossible to create file: %s', srtname)
|
||||
return None
|
||||
|
||||
|
||||
srt_length = fstat(srtfd).st_size
|
||||
if srt_length > 0:
|
||||
res.append((srtfd, lang))
|
||||
|
||||
return res
|
||||
+5
-162
@@ -58,6 +58,11 @@ from tscut.tools.timeframe import (
|
||||
parse_time_interval
|
||||
)
|
||||
from tscut.tools.discovery import check_required_tools
|
||||
from tscut.ocr.subtitles import (
|
||||
do_ocr,
|
||||
extract_srt,
|
||||
get_tesseract_supported_lang
|
||||
)
|
||||
|
||||
|
||||
# Useful SPS/PPS discussion.
|
||||
@@ -85,169 +90,7 @@ from tscut.tools.discovery import check_required_tools
|
||||
|
||||
|
||||
|
||||
@typechecked
|
||||
def get_tesseract_supported_lang(tesseract_path:str) -> dict[Lang, str]|None:
|
||||
"""
|
||||
Retrieves the set of natural languages supported by the Tesseract OCR tool.
|
||||
|
||||
This function runs the Tesseract binary with the --list-langs option and parses the output
|
||||
to extract the supported languages.
|
||||
|
||||
Args:
|
||||
tesseract_path (str): The path to the Tesseract binary.
|
||||
|
||||
Returns:
|
||||
dict[Lang, str] | None:
|
||||
- A dictionary mapping Lang objects to their corresponding language codes
|
||||
(e.g., "eng" for English)
|
||||
- None if an error occurs while running the Tesseract binary
|
||||
"""
|
||||
logger = logging.getLogger(__name__)
|
||||
res = {}
|
||||
|
||||
with Popen([tesseract_path, '--list-langs'], stdout=PIPE) as tesseract:
|
||||
for line in tesseract.stdout:
|
||||
line = line.decode('utf8')
|
||||
p = re.compile('(?P<lang>[a-z]{3})\n')
|
||||
m = re.match(p,line)
|
||||
if m is not None:
|
||||
try:
|
||||
lang = m.group('lang')
|
||||
key = Lang(lang)
|
||||
res[key] = lang
|
||||
except InvalidLanguageValue as e:
|
||||
logger.warning('Invalid language: %s', e)
|
||||
pass
|
||||
|
||||
tesseract.wait()
|
||||
|
||||
if tesseract.returncode != 0:
|
||||
logger.error("Tesseract returns an error code: %d",tesseract.returncode)
|
||||
return None
|
||||
|
||||
return res
|
||||
|
||||
|
||||
|
||||
@typechecked
|
||||
def extract_srt(mkvextract:str, filename:str, subtitles:dict[str, list[int]],
|
||||
langs:dict[Lang,str]) -> list[tuple[str,str,str,str]]|None:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
params = [mkvextract, filename, 'tracks']
|
||||
|
||||
res = []
|
||||
|
||||
for lang in subtitles:
|
||||
iso = Lang(lang)
|
||||
|
||||
if iso in langs:
|
||||
ocrlang = langs[iso]
|
||||
else:
|
||||
logger.warning("Language not supported by Tesseract: %s", iso.name)
|
||||
ocrlang ='osd'
|
||||
|
||||
if len(subtitles[lang]) == 1:
|
||||
params.append(f'{subtitles[lang][0]:d}:{lang}')
|
||||
res.append((f'{lang}.idx', f'{lang}.sub', lang, ocrlang))
|
||||
else:
|
||||
count = 1
|
||||
for track in subtitles[lang]:
|
||||
params.append(f'{track:d}:{lang}-{count:d}')
|
||||
res.append((f'{lang}-{count:d}.idx', f'{lang}-{count:d}.sub', lang, ocrlang))
|
||||
count = count+1
|
||||
|
||||
logger.debug('Executing %s', params)
|
||||
|
||||
env = {**os.environ, 'LANG': 'C'}
|
||||
with Popen(params, stdout=PIPE, close_fds=False, env=env) as extract:
|
||||
pb = tqdm(TextIOWrapper(extract.stdout, encoding="utf-8"), total=100, unit='%',
|
||||
desc='Extraction:')
|
||||
for line in pb:
|
||||
if line.startswith('Progress :'):
|
||||
p = re.compile('^Progress : (?P<progress>[0-9]{1,3})%$')
|
||||
m = p.match(line)
|
||||
if m is None:
|
||||
logger.error('Impossible to parse progress')
|
||||
pb.update(int(m['progress'])-pb.n)
|
||||
pb.update(100-pb.n)
|
||||
pb.refresh()
|
||||
pb.close()
|
||||
|
||||
extract.wait()
|
||||
|
||||
# mkvextract returns 0, 1 or 2 as error code.
|
||||
match extract.returncode:
|
||||
case 0:
|
||||
logger.info('Subtitle tracks were succesfully extracted.')
|
||||
case 1:
|
||||
logger.warning('Mkvextract returns warning')
|
||||
case 2:
|
||||
logger.error('Mkvextract returns an error code: %d', extract.returncode)
|
||||
res = None
|
||||
|
||||
return res
|
||||
|
||||
@typechecked
|
||||
def do_ocr(vobsubocr:str, idxs: list[tuple[str,str,str,str]], duration:timedelta,
|
||||
temporaries:list[IO[bytes]], dump_mem_fd:bool=False):
|
||||
logger = logging.getLogger(__name__)
|
||||
res = []
|
||||
|
||||
for idx_name, _, lang, iso in idxs:
|
||||
srtname = f'{os.path.splitext(idx_name)[0]}.srt'
|
||||
# Tesseract seems to recognize the three dots ... as "su"
|
||||
ldots = re.compile('^su\n$')
|
||||
# Timestamps produced by vobsubocr: 01:52:19,861 --> 01:52:21,641
|
||||
timestamps = re.compile((r'^[0-9]{2}:[0-9]{2}:[0-9]{2},[0-9]{3} \-\-> (?P<hours>[0-9]{2}):'
|
||||
r'(?P<minutes>[0-9]{2}):(?P<seconds>[0-9]{2}),[0-9]{3}$'))
|
||||
|
||||
srtfd = memfd_create(srtname, flags=0)
|
||||
with Popen([vobsubocr, '--lang', iso, idx_name], stdout=PIPE) as ocr:
|
||||
pb = tqdm(TextIOWrapper(ocr.stdout, encoding="utf-8"), total=
|
||||
int(duration/timedelta(seconds=1)), unit='s', desc='OCR')
|
||||
for line in pb:
|
||||
m = re.match(ldots,line)
|
||||
if m is not None:
|
||||
write(srtfd, '...'.encode(encoding='UTF-8'))
|
||||
else:
|
||||
write(srtfd, line.encode(encoding='UTF-8'))
|
||||
|
||||
m = re.match(timestamps, line)
|
||||
if m is not None:
|
||||
hours = int(m.group('hours'))
|
||||
minutes = int(m.group('hours'))
|
||||
seconds = int(m.group('seconds'))
|
||||
ts = timedelta(hours=hours, minutes=minutes, seconds=seconds)
|
||||
pb.n = int(ts/timedelta(seconds=1))
|
||||
pb.update()
|
||||
|
||||
status = ocr.wait()
|
||||
|
||||
if status != 0:
|
||||
logger.error('OCR failed with status code: %d', status)
|
||||
|
||||
if dump_mem_fd:
|
||||
try:
|
||||
with open(srtname,'w', encoding='utf8') as dump_srt:
|
||||
lseek(srtfd, 0, SEEK_SET)
|
||||
srt_length = fstat(srtfd).st_size
|
||||
buf = read(srtfd, srt_length)
|
||||
outfd = dump_srt.fileno()
|
||||
pos = 0
|
||||
while pos < srt_length:
|
||||
pos+=write(outfd, buf[pos:])
|
||||
temporaries.append(dump_srt)
|
||||
except OSError:
|
||||
logger.error('Impossible to create file: %s', srtname)
|
||||
return None
|
||||
|
||||
|
||||
srt_length = fstat(srtfd).st_size
|
||||
if srt_length > 0:
|
||||
res.append((srtfd, lang))
|
||||
|
||||
return res
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user