ocr: run dir_in mode in parallel (like layout), too…

- add CLI options `--num-jobs` and `--halt-fail`
- separate `Eynollah_ocr.run_single()` to be scheduled
- use ProcessPoolExecutor w/ forking and QueueListener
- also skip if input XML file is missing
- log processing times per job and overall

src/eynollah/cli/cli_ocr.py

@@ -47,6 +47,20 @@ import click
     help="overwrite (instead of skipping) if output xml exists",
     is_flag=True,
 )
+@click.option(
+    "--num-jobs",
+    "-j",
+    default=0,
+    type=click.IntRange(min=0),
+    help="number of parallel images to process (for --dir_in mode; also helps better utilise GPU if available); 0 means based on autodetected number of processor cores",
+)
+@click.option(
+    "--halt-fail",
+    "-H",
+    default=0,
+    type=click.FloatRange(min=0),
+    help="abort when number of failed images exceeds this value (if >=1) or ratio of failed over total images exceeds this value (if <1); 0 means ignore failures",
+)
 @click.option(
     "--tr_ocr",
     "-trocr",
@@ -83,6 +97,8 @@ def ocr_cli(
     out,
     dir_out_image_text,
     overwrite,
+    num_jobs,
+    halt_fail,
     tr_ocr,
     do_not_mask_with_textline_contour,
     batch_size,
@@ -108,4 +124,6 @@ def ocr_cli(
                      dir_xmls=dir_xmls,
                      dir_out_image_text=dir_out_image_text,
                      dir_out=out,
+                     num_jobs=num_jobs,
+                     halt_fail=halt_fail,
     )

src/eynollah/eynollah_ocr.py

@@ -1,13 +1,17 @@
 # FIXME: fix all of those...
 # pyright: reportOptionalSubscript=false
 
-from logging import Logger, getLogger
+import logging
+import logging.handlers
 from typing import List, Optional
 from pathlib import Path
 import os
 import gc
 import math
+import time
 from dataclasses import dataclass
+import multiprocessing as mp
+from concurrent.futures import ProcessPoolExecutor, as_completed
 
 import cv2
 from cv2.typing import MatLike
@@ -37,6 +41,21 @@ from .utils.utils_ocr import (
     rotate_image_with_padding,
 )
 
+
+_instance = None
+def _set_instance(instance):
+    global _instance
+    _instance = instance
+def _run_single(*args, **kwargs):
+    logq = kwargs.pop('logq')
+    # replace all inherited handlers with queue handler
+    logging.root.handlers.clear()
+    _instance.logger.handlers.clear()
+    handler = logging.handlers.QueueHandler(logq)
+    logging.root.addHandler(handler)
+    return _instance.run_single(*args, **kwargs)
+
+
 # TODO: refine typing
 @dataclass
 class EynollahOcrResult:
@@ -54,13 +73,13 @@ class Eynollah_ocr(Eynollah):
         batch_size: int=0,
         do_not_mask_with_textline_contour: bool=False,
         min_conf_value_of_textline_text : float=0.3,
-        logger: Optional[Logger]=None,
+        logger: Optional[logging.Logger]=None,
         device: str = '',
     ):
         self.tr_ocr = tr_ocr
         # masking for OCR and GT generation, relevant for skewed lines and bounding boxes
         self.do_not_mask_with_textline_contour = do_not_mask_with_textline_contour
-        self.logger = logger if logger else getLogger('eynollah.ocr')
+        self.logger = logger if logger else logging.getLogger('eynollah.ocr')
         
         self.min_conf_value_of_textline_text = min_conf_value_of_textline_text
         self.b_s = batch_size or 2 if tr_ocr else 8
@@ -416,16 +435,17 @@ class Eynollah_ocr(Eynollah):
         self.logger.info("output filename: '%s'", out_file_ocr)
         page_tree.write(out_file_ocr, xml_declaration=True, method='xml', encoding="utf-8", default_namespace=None)
 
-    def run(
+    def run(self,
-        self,
+            *,
-        *,
+            overwrite: bool = False,
-        overwrite: bool = False,
+            dir_in: str = "",
-        dir_in: Optional[str] = None,
+            dir_in_bin: str = "",
-        dir_in_bin: Optional[str] = None,
+            image_filename: str = "",
-        image_filename: Optional[str] = None,
+            dir_xmls: str,
-        dir_xmls: str,
+            dir_out_image_text: str = "",
-        dir_out_image_text: Optional[str] = None,
+            dir_out: str,
-        dir_out: str,
+            num_jobs: int = 0,
+            halt_fail: float = 0,
     ):
         """
         Run OCR.
@@ -435,58 +455,115 @@ class Eynollah_ocr(Eynollah):
             dir_in_bin (str): Prediction with RGB and binarized images for selected pages, should not be the default
         """
         if dir_in:
+            t0_tot = time.time()
             ls_imgs = [os.path.join(dir_in, image_filename)
-                    for image_filename in filter(is_image_filename,
+                       for image_filename in filter(is_image_filename,
                                                     os.listdir(dir_in))]
+            with ProcessPoolExecutor(max_workers=num_jobs or None,
+                                     mp_context=mp.get_context('fork'),
+                                     initializer=_set_instance,
+                                     initargs=(self,)
+            ) as exe:
+                jobs = {}
+                mngr = mp.get_context('fork').Manager()
+                n_success = n_fail = 0
+                for img_filename in ls_imgs:
+                    logq = mngr.Queue()
+                    jobs[exe.submit(_run_single, img_filename,
+                                    dir_out=dir_out,
+                                    dir_xmls=dir_xmls,
+                                    dir_in_bin=dir_in_bin,
+                                    dir_out_image_text=dir_out_image_text,
+                                    overwrite=overwrite,
+                                    logq=logq)] = img_filename, logq
+                for job in as_completed(list(jobs)):
+                    img_filename, logq = jobs[job]
+                    loglistener = logging.handlers.QueueListener(
+                        logq, *self.logger.handlers, respect_handler_level=False)
+                    try:
+                        loglistener.start()
+                        job.result()
+                        n_success += 1
+                    except:
+                        self.logger.exception("Job %s failed", img_filename)
+                        n_fail += 1
+                        if (halt_fail and
+                            n_fail >= halt_fail * (len(jobs) if halt_fail < 1 else 1)):
+                            self.logger.fatal("terminating after %d failures", n_fail)
+                            for job in jobs:
+                                job.cancel()
+                            break
+                    finally:
+                        loglistener.stop()
+            self.logger.info("%d of %d jobs successful", n_success, len(jobs))
+            self.logger.info("All jobs done in %.1fs", time.time() - t0_tot)
         else:
             assert image_filename
-            ls_imgs = [image_filename]
+            self.run_single(image_filename,
+                            dir_xmls=dir_xmls,
+                            dir_out=dir_out,
+                            dir_in_bin=dir_in_bin,
+                            dir_out_image_text=dir_out_image_text,
+                            overwrite=overwrite)
 
-        for img_filename in ls_imgs:
+    def run_single(self, 
-            file_stem = Path(img_filename).stem
+                   img_filename: str,
-            page_file_in = os.path.join(dir_xmls, file_stem+'.xml')
+                   dir_xmls: str,
-            out_file_ocr = os.path.join(dir_out, file_stem+'.xml')
+                   dir_out: str = "",
-            
+                   dir_in_bin: str = "",
-            if os.path.exists(out_file_ocr):
+                   dir_out_image_text: str = "",
-                if overwrite:
+                   overwrite: bool = False,
-                    self.logger.warning("will overwrite existing output file '%s'", out_file_ocr)
+    ):
-                else:
+        file_stem = Path(img_filename).stem
-                    self.logger.warning("will skip input for existing output file '%s'", out_file_ocr)
+        page_file_in = os.path.join(dir_xmls, file_stem + '.xml')
-                    continue
+        out_file_ocr = os.path.join(dir_out, file_stem + '.xml')
-                
-            img = cv2.imread(img_filename)
-            self.logger.info(img_filename)
-            page_tree = ET.parse(page_file_in, parser = ET.XMLParser(encoding="utf-8"))
-            page_ns = etree_namespace_for_element_tag(page_tree.getroot().tag)
 
-            out_image_with_text = None
+        if os.path.exists(out_file_ocr):
-            if dir_out_image_text:
+            if overwrite:
-                out_image_with_text = os.path.join(dir_out_image_text, file_stem + '.png')
+                self.logger.warning("will overwrite existing output file '%s'", out_file_ocr)
-
-            img_bin = None
-            if dir_in_bin:
-                img_bin = cv2.imread(os.path.join(dir_in_bin, file_stem+'.png'))
-
-
-            if self.tr_ocr:
-                result = self.run_trocr(
-                    img=img,
-                    page_tree=page_tree,
-                    page_ns=page_ns,
-                )
             else:
-                result = self.run_cnn( 
+                self.logger.warning("will skip input for existing output file '%s'", out_file_ocr)
-                    img=img,
+                return
-                    page_tree=page_tree,
+        if not os.path.exists(page_file_in):
-                    page_ns=page_ns,
+            self.logger.error("will skip missing input file '%s'", page_file_in)
-                    img_bin=img_bin,
+            return
-                )
 
-            self.write_ocr(
+        t0 = time.time()
-                result=result,
+
+        img = cv2.imread(img_filename)
+        self.logger.info(img_filename)
+        page_tree = ET.parse(page_file_in, parser = ET.XMLParser(encoding="utf-8"))
+        page_ns = etree_namespace_for_element_tag(page_tree.getroot().tag)
+
+        out_image_with_text = None
+        if dir_out_image_text:
+            out_image_with_text = os.path.join(dir_out_image_text, file_stem + '.png')
+
+        img_bin = None
+        if dir_in_bin:
+            img_bin = cv2.imread(os.path.join(dir_in_bin, file_stem+'.png'))
+
+
+        if self.tr_ocr:
+            result = self.run_trocr(
                 img=img,
                 page_tree=page_tree,
                 page_ns=page_ns,
-                out_file_ocr=out_file_ocr,
-                out_image_with_text=out_image_with_text,
             )
+        else:
+            result = self.run_cnn( 
+                img=img,
+                page_tree=page_tree,
+                page_ns=page_ns,
+                img_bin=img_bin,
+            )
+
+        self.write_ocr(
+            result=result,
+            img=img,
+            page_tree=page_tree,
+            page_ns=page_ns,
+            out_file_ocr=out_file_ocr,
+            out_image_with_text=out_image_with_text,
+        )
+        self.logger.info("Job done in %.1fs", time.time() - t0)