Shared im_mode setting in GroundTruthDataset/PolygonGTDataset

mittagessen · Apr 8, 2024 · 5524d88 · 5524d88
1 parent 46be12c
commit 5524d88
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 12 deletions.
diff --git 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
@@ -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: