diff --git a/src/tscut/ocr/subtitles.py b/src/tscut/ocr/subtitles.py new file mode 100644 index 0000000..f62076f --- /dev/null +++ b/src/tscut/ocr/subtitles.py @@ -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[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[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[0-9]{2}):' + r'(?P[0-9]{2}):(?P[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 diff --git a/src/tscut/tscut.py b/src/tscut/tscut.py index fd98031..9b4a5a0 100755 --- a/src/tscut/tscut.py +++ b/src/tscut/tscut.py @@ -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[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[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[0-9]{2}):' - r'(?P[0-9]{2}):(?P[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