diff --git a/kraken/lib/dataset/recognition.py b/kraken/lib/dataset/recognition.py index 19516260b..5a13d2f00 100644 --- a/kraken/lib/dataset/recognition.py +++ b/kraken/lib/dataset/recognition.py @@ -30,8 +30,9 @@ Tuple, Union) from PIL import Image -from torch.utils.data import Dataset +from ctypes import c_char from torchvision import transforms +from torch.utils.data import Dataset from kraken.containers import BaselineLine, BBoxLine, Segmentation from kraken.lib import functional_im_transforms as F_t @@ -333,7 +334,7 @@ def __init__(self, if augmentation: self.aug = DefaultAugmenter() - self.im_mode = mp.Value(c_char, b'1') + self._im_mode = mp.Value(c_char, b'1') def add(self, line: Optional[BaselineLine] = None, @@ -445,10 +446,10 @@ def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]: if is_bitonal(im): im_mode = b'1' - with self.im_mode.get_lock(): - if im_mode > self.im_mode: - logger.info(f'Upgrading "im_mode" from {self.im_mode} to {im_mode}') - self.im_mode = im_mode + with self._im_mode.get_lock(): + if im_mode > self._im_mode.value: + logger.info(f'Upgrading "im_mode" from {self._im_mode.value} to {im_mode}') + self._im_mode.value = im_mode if self.aug: im = im.permute((1, 2, 0)).numpy() o = self.aug(image=im) @@ -464,6 +465,12 @@ def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]: def __len__(self) -> int: return len(self._images) + @property + def im_mode(self): + return {b'1': '1', + b'L': 'L', + b'R': 'RGB'}[self._im_mode.value] + class GroundTruthDataset(Dataset): """ @@ -523,7 +530,7 @@ def __init__(self, if augmentation: self.aug = DefaultAugmenter() - self.im_mode = mp.Value(c_char, b'1') + self._im_mode = mp.Value(c_char, b'1') def add(self, line: Optional[BBoxLine] = None, @@ -624,10 +631,10 @@ def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]: im_mode = b'L' if is_bitonal(im): im_mode = b'1' - with self.im_mode.get_lock(): - if im_mode > self.im_mode: - logger.info(f'Upgrading "im_mode" from {self.im_mode} to {im_mode}') - self.im_mode = im_mode + with self._im_mode.get_lock(): + if im_mode > self._im_mode.value: + logger.info(f'Upgrading "im_mode" from {self._im_mode.value} to {im_mode}') + self._im_mode.value = im_mode if self.aug: im = im.permute((1, 2, 0)).numpy() o = self.aug(image=im) @@ -643,3 +650,11 @@ def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]: def __len__(self) -> int: return len(self._images) + + @property + def im_mode(self): + return {b'1': '1', + b'L': 'L', + b'R': 'RGB'}[self._im_mode.value] + + diff --git a/kraken/lib/util.py b/kraken/lib/util.py index 609e924fe..267b4545b 100644 --- a/kraken/lib/util.py +++ b/kraken/lib/util.py @@ -56,7 +56,7 @@ def is_bitonal(im: Union[Image.Image, torch.Tensor]) -> bool: if isinstance(im, Image.Image): return im.getcolors(2) is not None and len(im.getcolors(2)) == 2 elif isinstance(im, torch.Tensor): - return len(im.int().unique()) == 2 + return len(im.unique()) == 2 def get_im_str(im: Image.Image) -> str: