Move functions related to subtitles OCR in a dedicated module.

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<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

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<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