diff --git a/src/eynollah/ocrd_cli_binarization.py b/src/eynollah/ocrd_cli_binarization.py index 848bbac..e5fa740 100644 --- a/src/eynollah/ocrd_cli_binarization.py +++ b/src/eynollah/ocrd_cli_binarization.py @@ -10,16 +10,7 @@ from ocrd_models.ocrd_page import OcrdPage, AlternativeImageType from ocrd.decorators import ocrd_cli_options, ocrd_cli_wrap_processor from .sbb_binarize import SbbBinarizer - - -def cv2pil(img): - return Image.fromarray(img.astype('uint8')) - -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) +from .utils.pil_cv2 import cv2pil, pil2cv class SbbBinarizeProcessor(Processor): # already employs GPU (without singleton process atm) diff --git a/src/eynollah/utils/pil_cv2.py b/src/eynollah/utils/pil_cv2.py index 9f6913e..7975939 100644 --- a/src/eynollah/utils/pil_cv2.py +++ b/src/eynollah/utils/pil_cv2.py @@ -11,8 +11,12 @@ def cv2pil(img): def pil2cv(img): # from ocrd/workspace.py - color_conversion = COLOR_GRAY2BGR if img.mode in ('1', 'L') else COLOR_RGB2BGR + color_conversion = COLOR_GRAY2BGR if img.mode in ('1', 'L', 'LA') else COLOR_RGB2BGR pil_as_np_array = np.array(img).astype('uint8') if img.mode == '1' else np.array(img) + if pil_as_np_array.shape[-1] == 2: + pil_as_np_array = pil_as_np_array[:,:,0] + elif pil_as_np_array.shape[-1] == 4: + pil_as_np_array = pil_as_np_array[:,:,:3] return cvtColor(pil_as_np_array, color_conversion) def check_dpi(img):