TALIE-1519: unify colors for images and pdfs

dedoc/attachments_extractors/concrete_attachments_extractors/image_attachments_extractor.py

@@ -70,7 +70,6 @@ def extract(self, file_path: str, parameters: Optional[dict] = None) -> List[Att
         attachments = []
 
         image = cv2.imread(file_path)
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         predictions = self._predict(image)
 
         for prediction in predictions:
@@ -86,7 +85,7 @@ def extract(self, file_path: str, parameters: Optional[dict] = None) -> List[Att
 
                 tmp_file_name = get_unique_name(filename)
                 tmp_file_path = os.path.join(attachments_dir, tmp_file_name)
-                cv2.imwrite(tmp_file_path, cv2.cvtColor(part, cv2.COLOR_RGB2BGR))
+                cv2.imwrite(tmp_file_path, part)
 
                 image_attachment = PdfImageAttachment(
                     original_name=tmp_file_name,

dedoc/attachments_extractors/concrete_attachments_extractors/pdf_attachments_extractor.py

@@ -88,7 +88,10 @@ def __get_root_attachments(self, reader: PdfReader) -> List[Tuple[str, bytes]]:
         import uuid
 
         attachments = []
-        catalog = reader.trailer["/Root"]
+        catalog = reader.trailer.get("/Root")
+        if catalog is None:
+            return attachments
+
         if "/Names" in catalog.keys() and "/EmbeddedFiles" in catalog["/Names"].keys() and "/Names" in catalog["/Names"]["/EmbeddedFiles"].keys():
             file_names = catalog["/Names"]["/EmbeddedFiles"]["/Names"]
             for f in file_names:

dedoc/readers/pdf_reader/pdf_auto_reader/txtlayer_detector.py

@@ -2,7 +2,7 @@
 import math
 from copy import deepcopy
 from itertools import chain
-from typing import List
+from typing import List, Optional
 
 import numpy as np
 
@@ -37,22 +37,24 @@ def detect_txtlayer(self, path: str, parameters: dict) -> List[TxtLayerResult]:
         if txtlayer_classifier is None:
             raise ValueError(f"Unknown textual layer classifier `{classifier_name}`")
 
+        start, end = get_param_page_slice(parameters)
+        start = 1 if start is None else start + 1
+
         classify_each_page = get_bool_parameter(parameters, "each_page_textual_layer_detection", False)
         detect_function = self.__classify_each_page if classify_each_page else self.__classify_all_pages
         try:
-            return detect_function(path, parameters, txtlayer_classifier)
+            return detect_function(path, parameters, txtlayer_classifier, start, end)
         except Exception as e:
             self.logger.debug(f"Error occurred white detecting PDF textual layer ({e})")
-            return [TxtLayerResult(correct=False, start=1, end=None)]
+            return [TxtLayerResult(correct=False, start=start, end=end)]
 
-    def __classify_all_pages(self, path: str, parameters: dict, txtlayer_classifier: AbstractTxtlayerClassifier) -> List[TxtLayerResult]:
+    def __classify_all_pages(
+            self, path: str, parameters: dict, txtlayer_classifier: AbstractTxtlayerClassifier, start: int, end: Optional[int]
+    ) -> List[TxtLayerResult]:
         """
         Check only first 8 pages of the document, use classification results for the entire document.
         Separately handle the first page (it's common that only first page doesn't have a textual layer).
         """
-        start, end = get_param_page_slice(parameters)
-        start = 1 if start is None else start + 1
-
         parameters_copy = deepcopy(parameters)
         parameters_copy["pages"] = "1:8"  # two batches for pdf_txtlayer_reader
         parameters_copy["need_pdf_table_analysis"] = "false"
@@ -72,13 +74,13 @@ def __classify_all_pages(self, path: str, parameters: dict, txtlayer_classifier:
         else:
             return [TxtLayerResult(correct=False, start=start, end=start), TxtLayerResult(correct=True, start=start + 1, end=end)]
 
-    def __classify_each_page(self, path: str, parameters: dict, txtlayer_classifier: AbstractTxtlayerClassifier) -> List[TxtLayerResult]:
+    def __classify_each_page(
+            self, path: str, parameters: dict, txtlayer_classifier: AbstractTxtlayerClassifier, start: int, end: Optional[int]
+    ) -> List[TxtLayerResult]:
         """
         Classify each page of the document correct/not correct textual layer.
         """
         document = self.pdf_reader.read(path, parameters=parameters)
-        start, end = get_param_page_slice(parameters)
-        start = 1 if start is None else start + 1
         if not document.lines:
             return [TxtLayerResult(correct=False, start=start, end=end)]
 

dedoc/readers/pdf_reader/pdf_base_reader.py

@@ -240,6 +240,7 @@ def _split_pdf2image(self, path: str, page_from: int, page_to: int) -> Iterator[
         if page_from >= page_to:
             return
 
+        import cv2
         import math
         import os
         import numpy as np
@@ -263,7 +264,8 @@ def _split_pdf2image(self, path: str, page_from: int, page_to: int) -> Iterator[
                     left += 1
                     if left > page_to + 1:
                         break
-                    yield np.array(image)
+                    image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)
+                    yield image
         except (PDFPageCountError, PDFSyntaxError) as error:
             raise BadFileFormatError(f"Bad pdf file:\n file_name = {os.path.basename(path)} \n exception = {error.args}")
 

labeling/train_dataset/taskers/images_creators/concrete_creators/docx_images_creator.py

@@ -11,6 +11,7 @@
 from typing import Dict, Iterable, Iterator, Optional, Tuple
 from typing import List
 
+import cv2
 import numpy as np
 from PIL import Image
 from PIL import ImageColor
@@ -214,7 +215,8 @@ def _create_images_from_pdf(self, pdfs: PairedPdf, page: List[dict], tmp_dir: st
         uid2path = defaultdict(list)
         n = 0
         for two_color, many_color in zip(two_color_images, many_color_images):
-
+            two_color = cv2.cvtColor(two_color, cv2.COLOR_RGB2BGR)
+            many_color = cv2.cvtColor(many_color, cv2.COLOR_RGB2BGR)
             diff = many_color - two_color
             all_masks = np.abs(diff) > 0
             many_color[all_masks] = 255