mirror of
https://github.com/qurator-spk/eynollah.git
synced 2026-07-26 05:29:16 +02:00
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
This commit is contained in:
Robert Sachunsky
parent
1c8ac38d31
commit
6840b67961
2 changed files with 151 additions and 56 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
self,
|
||||
def run(self,
|
||||
*,
|
||||
overwrite: bool = False,
|
||||
dir_in: Optional[str] = None,
|
||||
dir_in_bin: Optional[str] = None,
|
||||
image_filename: Optional[str] = None,
|
||||
dir_in: str = "",
|
||||
dir_in_bin: str = "",
|
||||
image_filename: str = "",
|
||||
dir_xmls: str,
|
||||
dir_out_image_text: Optional[str] = None,
|
||||
dir_out_image_text: str = "",
|
||||
dir_out: str,
|
||||
num_jobs: int = 0,
|
||||
halt_fail: float = 0,
|
||||
):
|
||||
"""
|
||||
Run OCR.
|
||||
|
|
@ -435,24 +455,80 @@ 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,
|
||||
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,
|
||||
img_filename: str,
|
||||
dir_xmls: str,
|
||||
dir_out: str = "",
|
||||
dir_in_bin: str = "",
|
||||
dir_out_image_text: str = "",
|
||||
overwrite: bool = False,
|
||||
):
|
||||
file_stem = Path(img_filename).stem
|
||||
page_file_in = os.path.join(dir_xmls, file_stem+'.xml')
|
||||
out_file_ocr = os.path.join(dir_out, file_stem+'.xml')
|
||||
page_file_in = os.path.join(dir_xmls, file_stem + '.xml')
|
||||
out_file_ocr = os.path.join(dir_out, file_stem + '.xml')
|
||||
|
||||
if os.path.exists(out_file_ocr):
|
||||
if overwrite:
|
||||
self.logger.warning("will overwrite existing output file '%s'", out_file_ocr)
|
||||
else:
|
||||
self.logger.warning("will skip input for existing output file '%s'", out_file_ocr)
|
||||
continue
|
||||
return
|
||||
if not os.path.exists(page_file_in):
|
||||
self.logger.error("will skip missing input file '%s'", page_file_in)
|
||||
return
|
||||
|
||||
t0 = time.time()
|
||||
|
||||
img = cv2.imread(img_filename)
|
||||
self.logger.info(img_filename)
|
||||
|
|
@ -490,3 +566,4 @@ class Eynollah_ocr(Eynollah):
|
|||
out_file_ocr=out_file_ocr,
|
||||
out_image_with_text=out_image_with_text,
|
||||
)
|
||||
self.logger.info("Job done in %.1fs", time.time() - t0)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue