segment/recognize: spawn bg process for Kraken…

- during `setup`, instead of loading models in the processor
  directly, instantiate and spawn a singleton predictor subprocess
  with the given parameters (after resolving the model path name),
  communicating via shared (task and result) queues to synchronize
  processor and predictor processes;
  the predictor will then load models in its own address space
- at runtime, the processor merely calls the predictor with the
  respective arguments for that page, which translates into
  - putting the arguments on the task queue
  - getting the results from the result queue, blocking
- at runtime, the predictor loops into:
  - receiving inputs from the task queue, blocking
  - calling `predict` on them
  - putting outputs on the result queue
- in the predictor, tasks and results are identified via page id,
  so results get retrieved for their respective task only,
  implemented via shared dict to synchronize forked processor workers
- during `shutdown`, tell the predictor to shut down as well
  (terminating the subprocess);
  the predictor will then exit its loop and close the queues
- abstract from kraken.pageseg, kraken.blla, and kraken.rpred
  differences in initialization phase and inference phase via
  shared `common.KrakenPredictor` class, override specifics in
  - `recognize.KrakenRecognizePredictor`:
    - during `setup`, after loading the model, submit a special "task"
      to query the model's `one_channel_mode` attribute
    - at runtime, translate the model into a `defaultdict` for `mm_rpred`,
      but picklable to be compatible with mp.Queue; for the same reason,
      exhaust the result generator immediately
  - `segment.KrakenSegmentPredictor`: during `setup`, map the given
    parameters and inputs to kwargs as applicable by either `pageseg.segment`
    or `blla.segment`

Author: bertsky

ocrd_kraken/binarize.py

@@ -2,10 +2,11 @@
 from os.path import join
 from typing import Optional
 
+import kraken.binarization
+
 from ocrd.processor.base import OcrdPageResult
 from ocrd.processor.ocrd_page_result import OcrdPageResultImage
 
-import kraken.binarization
 from ocrd import Processor
 from ocrd_utils import assert_file_grp_cardinality, getLogger, make_file_id, MIMETYPE_PAGE
 from ocrd_models.ocrd_page import AlternativeImageType, OcrdPage, to_xml

ocrd_kraken/common.py

@@ -0,0 +1,76 @@
+import multiprocessing as mp
+
+from ocrd_utils import config, initLogging
+
+class KrakenPredictor(mp.context.SpawnProcess):
+    def __init__(self, logger, parameter):
+        self.logger = logger
+        self.parameter = parameter
+        ctxt = mp.get_context('spawn')
+        self.taskq = ctxt.Queue(maxsize=1 + config.OCRD_MAX_PARALLEL_PAGES)
+        self.resultq = ctxt.Queue(maxsize=1 + config.OCRD_MAX_PARALLEL_PAGES)
+        self.terminate = ctxt.Event()
+        ctxt = mp.get_context('fork') # base.Processor will fork workers
+        self.results = ctxt.Manager().dict()
+        super().__init__()
+        self.daemon = True
+    def __call__(self, page_id, *page_input):
+        self.taskq.put((page_id, page_input))
+        self.logger.debug("sent task for '%s'", page_id)
+        #return self.get(page_id)
+        result = self.get(page_id)
+        self.logger.debug("received result for '%s'", page_id)
+        return result
+    def get(self, page_id):
+        while not self.terminate.is_set():
+            if page_id in self.results:
+                result = self.results.pop(page_id)
+                if isinstance(result, Exception):
+                    raise Exception(f"predictor failed for {page_id}") from result
+                return result
+            try:
+                page_id, result = self.resultq.get(timeout=0.7)
+            except mp.queues.Empty:
+                continue
+            self.logger.debug("storing results for '%s'", page_id)
+            self.results[page_id] = result
+        raise Exception(f"predictor terminated while waiting on results for {page_id}")
+    def run(self):
+        initLogging()
+        try:
+            self.setup()
+        except Exception as e:
+            self.logger.exception("setup failed")
+            self.terminate.set()
+        while not self.terminate.is_set():
+            try:
+                page_id, page_input = self.taskq.get(timeout=1.1)
+            except mp.queues.Empty:
+                continue
+            self.logger.debug("predicting '%s'", page_id)
+            try:
+                page_output = self.predict(*page_input)
+            except Exception as e:
+                self.logger.error("prediction failed: %s", e.__class__.__name__)
+                page_output = e
+            self.resultq.put((page_id, page_output))
+            self.logger.debug("sent result for '%s'", page_id)
+        self.resultq.close()
+        self.resultq.cancel_join_thread()
+        self.logger.debug("predictor terminated")
+    def setup(self):
+        raise NotImplementedError()
+    def predict(self, *inputs):
+        raise NotImplementedError()
+    def shutdown(self):
+        # do not terminate from forked processor instances
+        if mp.parent_process() is None:
+            self.terminate.set()
+            self.taskq.close()
+            self.taskq.cancel_join_thread()
+            self.logger.debug(f"terminated {self} in {mp.current_process().name}")
+        else:
+            self.logger.debug(f"not touching {self} in {mp.current_process().name}")
+    def __del__(self):
+        self.logger.debug(f"deinit of {self} in {mp.current_process().name}")
+        self.shutdown()

ocrd_kraken/recognize.py

@@ -2,6 +2,7 @@
 from ocrd.processor.base import OcrdPageResult
 import regex
 import itertools
+from collections import defaultdict
 import numpy as np
 from scipy.sparse.csgraph import minimum_spanning_tree
 from shapely.geometry import Polygon, LineString, box as Rectangle
@@ -37,6 +38,38 @@
     TextLineOrderSimpleType
 )
 
+from .common import KrakenPredictor
+
+class KrakenRecognizePredictor(KrakenPredictor):
+    # workaround for Kraken's unpicklable defaultdict choice
+    class DefaultDict(defaultdict):
+        def __init__(self, default=None):
+            self.default = default
+            super().__init__()
+        def default_factory(self):
+            return self.default
+    def setup(self):
+        import torch
+        from kraken.lib.models import load_any
+        model = self.parameter['model']
+        self.logger.info("loading model '%s'", model)
+        device = self.parameter['device']
+        if device != 'cpu' and not torch.cuda.is_available():
+            device = 'cpu'
+        if device == 'cpu':
+            self.logger.warning("no CUDA device available. Running without GPU will be slow")
+        self.model = load_any(model, device=device)
+    def predict(self, *inputs):
+        from kraken.rpred import mm_rpred
+        if not len(inputs):
+            return self.model.nn.input[1] == 1 and self.model.one_channel_mode == '1'
+        image, segmentation = inputs
+        nets = __class__.DefaultDict(self.model)
+        result = mm_rpred(nets, image, segmentation,
+                          self.parameter['pad'],
+                          self.parameter['bidi_reordering'])
+        # we must exhaust the generator before enqueuing
+        return list(result)
 
 class KrakenRecognize(Processor):
 
@@ -48,23 +81,17 @@ def setup(self):
         """
         Load model, set predict function
         """
+        parameter = dict(self.parameter)
+        parameter['model'] = self.resolve_resource(parameter['model'])
+        self.predictor = KrakenRecognizePredictor(self.logger, parameter)
+        self.predictor.start()
+        self.binary = self.predictor("") # blocks until model is loaded
+        self.logger.info("loaded %s model %s", "binary" if self.binary else "grayscale", self.parameter["model"])
 
-        import torch
-        from kraken.rpred import rpred
-        from kraken.lib.models import load_any
-        model_fname = self.resolve_resource(self.parameter['model'])
-        self.logger.info("loading model '%s'", model_fname)
-        device = self.parameter['device']
-        if device != 'cpu' and not torch.cuda.is_available():
-            device = 'cpu'
-        if device == 'cpu':
-            self.logger.warning("no CUDA device available. Running without GPU will be slow")
-        self.model = load_any(model_fname, device=device)
-        def predict(page_image, segmentation):
-            return rpred(self.model, page_image, segmentation,
-                         self.parameter['pad'],
-                         self.parameter['bidi_reordering'])
-        self.predict = predict
+    def shutdown(self):
+        if getattr(self, 'predictor', None):
+            self.predictor.shutdown()
+            del self.predictor
 
     def process_page_pcgts(self, *input_pcgts: Optional[OcrdPage], page_id: Optional[str] = None) -> OcrdPageResult:
         """Recognize text on lines with Kraken.
@@ -96,8 +123,7 @@ def process_page_pcgts(self, *input_pcgts: Optional[OcrdPage], page_id: Optional
         page_image, page_coords, _ = self.workspace.image_from_page(
             page, page_id,
             feature_selector="binarized"
-            if self.model.nn.input[1] == 1 and self.model.one_channel_mode == '1'
-            else '')
+            if self.binary else '')
         page_rect = Rectangle(0, 0, page_image.width - 1, page_image.height - 1)
         # TODO: find out whether kraken.lib.xml.XMLPage(...).to_container() is adequate
 
@@ -152,7 +178,7 @@ def process_page_pcgts(self, *input_pcgts: Optional[OcrdPage], page_id: Optional
                                     text_direction='horizontal-lr',
                                     type=segtype,
                                     imagename=page_id)
-        for idx_line, ocr_record in enumerate(self.predict(page_image, segmentation)):
+        for idx_line, ocr_record in enumerate(self.predictor(page_id, page_image, segmentation)):
             line = all_lines[idx_line]
             id_line = line.id
             if not ocr_record.prediction and not ocr_record.cuts:

ocrd_kraken/segment.py

@@ -1,6 +1,10 @@
 from typing import Optional
 from PIL import ImageOps
 
+import shapely.geometry as geom
+from shapely.prepared import prep as geom_prep
+import torch
+
 from ocrd import Processor
 from ocrd.processor.ocrd_page_result import OcrdPageResult
 from ocrd_utils import (
@@ -22,44 +26,66 @@
     BaselineType,
 )
 
-import shapely.geometry as geom
-from shapely.prepared import prep as geom_prep
-import torch
-
-class KrakenSegment(Processor):
-
-    @property
-    def executable(self):
-        return 'ocrd-kraken-segment'
+from .common import KrakenPredictor
 
+class KrakenSegmentPredictor(KrakenPredictor):
     def setup(self):
-        """
-        Load models
-        """
-        kwargs = {}
-        kwargs['text_direction'] = self.parameter['text_direction']
-        self.use_legacy = self.parameter['use_legacy']
+        self.use_legacy = self.parameter.pop('use_legacy')
         if self.use_legacy:
-            from kraken.pageseg import segment
-            kwargs['scale'] = self.parameter['scale']
-            kwargs['maxcolseps'] = self.parameter['maxcolseps']
-            kwargs['black_colseps'] = self.parameter['black_colseps']
             self.logger.info("Using legacy segmenter")
+            # adapt to Kraken v5 changes:
+            self.parameter['no_hlines'] = self.parameter.pop('remove_hlines')
+            self.parameter.pop('device')
         else:
             from kraken.lib.vgsl import TorchVGSLModel
-            from kraken.blla import segment
             self.logger.info("Using blla segmenter")
-            blla_model_fname = self.resolve_resource(self.parameter['blla_model'])
-            kwargs['model'] = TorchVGSLModel.load_model(blla_model_fname)
+            self.logger.info("loading model '%s'", self.parameter['model'])
+            self.parameter['model'] = TorchVGSLModel.load_model(self.parameter['model'])
             device = self.parameter['device']
             if device != 'cpu' and not torch.cuda.is_available():
                 device = 'cpu'
             if device == 'cpu':
                 self.logger.warning("no CUDA device available. Running without GPU will be slow")
-            kwargs['device'] = device
-        def segmenter(img, mask=None):
-            return segment(img, mask=mask, **kwargs)
-        self.segmenter = segmenter
+            self.parameter['device'] = device
+            # adapt to Kraken v5 changes:
+            self.parameter.pop('scale')
+            self.parameter.pop('remove_hlines')
+            self.parameter.pop('maxcolseps')
+            self.parameter.pop('black_colseps')
+    def predict(self, *inputs):
+        if self.use_legacy:
+            from kraken.pageseg import segment
+        else:
+            from kraken.blla import segment
+        image, mask = inputs
+        return segment(image, mask=mask, **self.parameter)
+
+class KrakenSegment(Processor):
+
+    @property
+    def executable(self):
+        return 'ocrd-kraken-segment'
+
+    def setup(self):
+        """
+        Load models
+        """
+        parameter = dict(self.parameter)
+        model = parameter.pop('blla_model')
+        del parameter['blla_classes']
+        del parameter['overwrite_segments']
+        del parameter['level-of-operation']
+        self.use_legacy = parameter['use_legacy']
+        if not self.use_legacy:
+            parameter['model'] = self.resolve_resource(model)
+        self.predictor = KrakenSegmentPredictor(self.logger, parameter)
+        self.predictor.start()
+
+    def shutdown(self):
+        import multiprocessing as mp
+        if getattr(self, 'predictor', None):
+            self.predictor.shutdown()
+            del self.predictor
 
     def process_page_pcgts(self, *input_pcgts: Optional[OcrdPage], page_id: Optional[str] = None) -> OcrdPageResult:
         """Segment into (regions and) lines with Kraken.
@@ -109,7 +135,7 @@ def process_page_pcgts(self, *input_pcgts: Optional[OcrdPage], page_id: Optional
                 page.TextRegion = []
             elif len(page.TextRegion or []):
                 self.logger.warning('Keeping %d text regions on page "%s"', len(page.TextRegion or []), page.id)
-            self._process_page(page_image, page_coords, page, zoom)
+            self._process_page(page_image, page_coords, page, page_id, zoom)
         elif self.parameter['level-of-operation'] == 'table':
             regions = page.get_AllRegions(classes=['Table'])
             if not regions:
@@ -120,7 +146,7 @@ def process_page_pcgts(self, *input_pcgts: Optional[OcrdPage], page_id: Optional
                     region.TextRegion = []
                 elif len(region.TextRegion or []):
                     self.logger.warning('Keeping %d text regions in region "%s"', len(region.TextRegion or []), region.id)
-                self._process_page(page_image, page_coords, region, zoom)
+                self._process_page(page_image, page_coords, region, page_id, zoom)
         else:
             regions = page.get_AllRegions(classes=['Text'])
             if not regions:
@@ -131,11 +157,11 @@ def process_page_pcgts(self, *input_pcgts: Optional[OcrdPage], page_id: Optional
                     region.TextLine = []
                 elif len(region.TextLine or []):
                     self.logger.warning('Keeping %d lines in region "%s"', len(region.TextLine or []), region.id)
-                self._process_region(page_image, page_coords, region, zoom)
+                self._process_region(page_image, page_coords, region, page_id, zoom)
 
         return OcrdPageResult(pcgts)
 
-    def _process_page(self, page_image, page_coords, page, zoom=1.0):
+    def _process_page(self, page_image, page_coords, page, page_id, zoom=1.0):
         def getmask():
             # use mask if existing regions (any type for page, text cells for table)
             # or segment is lower than page level
@@ -173,10 +199,10 @@ def getmask():
                 # poly = geom.Polygon(poly).buffer(20/zoom).exterior.coords[:-1]
                 mask.paste(255, mask=polygon_mask(page_image, poly))
             return mask
-        res = self.segmenter(page_image, mask=getmask())
+        res = self.predictor(page_id, page_image, getmask())
         self.logger.debug("Finished segmentation, serializing")
+        #self.logger.debug(res)
         if self.use_legacy:
-            self.logger.debug(res)
             idx_line = 0
             for idx_line, line in enumerate(res.lines):
                 line_poly = polygon_from_x0y0x1y1(line.bbox)
@@ -191,7 +217,6 @@ def getmask():
                 page.add_TextRegion(region_elem)
             self.logger.debug("Found %d lines on page %s", idx_line + 1, page.id)
         else:
-            self.logger.debug(res)
             handled_lines = {}
             regions = [(type_, region)
                        for type_ in res.regions
@@ -245,7 +270,7 @@ def getmask():
                     page.add_TextRegion(region_elem)
             self.logger.debug("Found %d lines and %d regions on page %s", idx_line + 1, idx_region + 1, page.id)
 
-    def _process_region(self, page_image, page_coords, region, zoom=1.0):
+    def _process_region(self, page_image, page_coords, region, page_id, zoom=1.0):
         def getmask():
             poly = coordinates_of_segment(region, page_image, page_coords)
             poly = geom.Polygon(poly).buffer(20/zoom).exterior.coords[:-1]
@@ -256,7 +281,7 @@ def getmask():
                 # poly = geom.Polygon(poly).buffer(20/zoom).exterior.coords[:-1]
                 mask.paste(255, mask=polygon_mask(page_image, poly))
             return mask
-        res = self.segmenter(page_image, mask=getmask())
+        res = self.predictor(page_id, page_image, getmask())
         self.logger.debug("Finished segmentation, serializing")
         idx_line = 0
         if self.use_legacy: