convert between cv2 and pil, DRY binarizer call

sbb_binarize/ocrd_cli.py

@@ -4,7 +4,11 @@ import os.path
 from pkg_resources import resource_string
 from json import loads
 
+from PIL import Image
+import numpy as np
+import cv2
 from click import command
+
 from ocrd_utils import (
     getLogger,
     assert_file_grp_cardinality,
@@ -21,6 +25,16 @@ from .sbb_binarize import SbbBinarizer
 OCRD_TOOL = loads(resource_string(__name__, 'ocrd-tool.json').decode('utf8'))
 TOOL = 'ocrd-sbb-binarize'
 
+def cv2pil(img):
+    color_coverted = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
+    return Image.fromarray(color_coverted)
+
+def pil2cv(img):
+    # from ocrd/workspace.py
+    color_conversion = cv2.COLOR_GRAY2BGR if img.mode in ('1', 'L') else  cv2.COLOR_RGB2BGR
+    pil_as_np_array = np.array(img).astype('uint8') if img.mode == '1' else np.array(img)
+    return cv2.cvtColor(pil_as_np_array, color_conversion)
+
 class SbbBinarizeProcessor(Processor):
 
     def __init__(self, *args, **kwargs):
@@ -28,6 +42,14 @@ class SbbBinarizeProcessor(Processor):
         kwargs['version'] = OCRD_TOOL['version']
         super().__init__(*args, **kwargs)
 
+    def _run_binarizer(self, img):
+        return cv2pil(
+                SbbBinarizer(
+                    image=pil2cv(img),
+                    model=self.model_path,
+                    patches=self.use_patches,
+                    save=None).run())
+
     def process(self):
         """
         Binarize with sbb_binarization
@@ -37,8 +59,8 @@ class SbbBinarizeProcessor(Processor):
         assert_file_grp_cardinality(self.output_file_grp, 1)
 
         oplevel = self.parameter['operation_level']
-        use_patches = self.parameter['patches']
-        model_path = self.parameter['model']
+        self.use_patches = self.parameter['patches'] # pylint: disable=attribute-defined-outside-init
+        self.model_path = self.parameter['model'] # pylint: disable=attribute-defined-outside-init
 
         for n, input_file in enumerate(self.input_files):
             file_id = make_file_id(input_file, self.output_file_grp)
@@ -52,12 +74,7 @@ class SbbBinarizeProcessor(Processor):
             if oplevel == 'page':
                 LOG.info("Binarizing on 'page' level in page '%s'", page_id)
                 page_image, page_xywh, _ = self.workspace.image_from_page(page, page_id)
-                bin_image = SbbBinarizer(
-                    image=page_image,
-                    model=model_path,
-                    patches=use_patches,
-                    save=None
-                ).run()
+                bin_image = self._run_binarizer(page_image)
                 # update METS (add the image file):
                 bin_image_path = self.workspace.save_image_file(bin_image,
                         file_id + '.IMG-BIN',
@@ -74,12 +91,7 @@ class SbbBinarizeProcessor(Processor):
                     region_image, region_xywh = self.workspace.image_from_segment(region, page_image, page_xywh)
 
                     if oplevel == 'region':
-                        region_image_bin = SbbBinarizer(
-                            image=region_image,
-                            model=model_path,
-                            patches=use_patches,
-                            save=None
-                        ).run()
+                        region_image_bin = self._run_binarizer(region_image)
                         region_image_bin_path = self.workspace.save_image_file(
                                 region_image_bin,
                                 "%s_%s.IMG-BIN" % (file_id, region.id),
@@ -94,12 +106,7 @@ class SbbBinarizeProcessor(Processor):
                             LOG.warning("Page '%s' region '%s' contains no text lines", page_id, region.id)
                         for line in lines:
                             line_image, line_xywh = self.workspace.image_from_segment(line, page_image, page_xywh)
-                            line_image_bin = SbbBinarizer(
-                                image=line_image,
-                                model=model_path,
-                                patches=use_patches,
-                                save=None
-                            ).run()
+                            line_image_bin = self._run_binarizer(line_image)
                             line_image_bin_path = self.workspace.save_image_file(
                                     line_image_bin,
                                     "%s_%s_%s.IMG-BIN" % (file_id, region.id, line.id),

sbb_binarize/sbb_binarize.py

@@ -8,6 +8,7 @@ from os.path import join
 from warnings import catch_warnings, simplefilter
 
 import numpy as np
+from PIL import Image
 import cv2
 environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 stderr = sys.stderr
@@ -23,9 +24,10 @@ class SbbBinarizer:
 
     # TODO use True/False for patches
     def __init__(self, model, image=None, image_path=None, patches='false', save=None):
-        if not(image or image_path) or (image and image_path):
-            raise ValueError("Must pass either a PIL image or an image_path")
-        if image:
+        if (image is not None and image_path is not None) or \
+               (image is None and image_path is None):
+            raise ValueError("Must pass either a opencv2 image or an image_path")
+        if image is not None:
             self.image = image
         else:
             self.image = cv2.imread(self.image)