TLDR-850 added multipage into tabby reader

Author: oksidgy

dedoc/readers/pdf_reader/data_classes/tables/cell.py

@@ -3,10 +3,11 @@
 from dedocutils.data_structures import BBox
 
 from dedoc.data_structures.annotation import Annotation
+from dedoc.data_structures.cell_with_meta import CellWithMeta
 from dedoc.data_structures.line_with_meta import LineWithMeta
 
 
-class Cell:
+class Cell(CellWithMeta):
 
     @staticmethod
     def copy_from(cell: "Cell",
@@ -41,35 +42,27 @@ def shift(self, shift_x: int, shift_y: int, image_width: int, image_height: int)
         if self.con_coord:
             self.con_coord.shift(shift_x=shift_x, shift_y=shift_y)
 
-    def __init__(self,
-                 x_top_left: int,
-                 x_bottom_right: int,
-                 y_top_left: int,
-                 y_bottom_right: int,
-                 id_con: int = -1,
-                 lines: Optional[List[LineWithMeta]] = None,
-                 is_attribute: bool = False,
-                 is_attribute_required: bool = False,
-                 rotated_angle: int = 0,
-                 uid: str = None,
+    def __init__(self, x_top_left: int, x_bottom_right: int, y_top_left: int, y_bottom_right: int, id_con: int = -1, lines: Optional[List[LineWithMeta]] = None,
+                 is_attribute: bool = False, is_attribute_required: bool = False, rotated_angle: int = 0, uid: str = None,
                  contour_coord: Optional[BBox] = None) -> None:
+
         import uuid
 
         assert x_top_left <= x_bottom_right
         assert y_top_left <= y_bottom_right
+
+        self.lines = [] if lines is None else lines
+        super().__init__(lines)
+
         self.x_top_left = x_top_left
         self.x_bottom_right = x_bottom_right
         self.y_top_left = y_top_left
         self.y_bottom_right = y_bottom_right
         self.id_con = id_con
-        self.lines = [] if lines is None else lines
         self.is_attribute = is_attribute
         self.is_attribute_required = is_attribute_required
         self.rotated_angle = rotated_angle
         self.cell_uid = f"cell_{uuid.uuid1()}" if uid is None else uid
-        self.colspan = 1
-        self.rowspan = 1
-        self.invisible = False
         self.con_coord = contour_coord or BBox(0, 0, 0, 0)
 
     def __str__(self) -> str:

dedoc/readers/pdf_reader/data_classes/tables/location.py

@@ -10,6 +10,7 @@ def __init__(self, page_number: int, bbox: BBox, name: str = "", rotated_angle:
         self.page_number = page_number
         self.bbox = bbox
         self.name = name
+        # TODO put self.order (change LineWithLocation, PdfImageAttachment, ScanTable)
         self.rotated_angle = rotated_angle
 
     def shift(self, shift_x: int, shift_y: int) -> None:

dedoc/readers/pdf_reader/data_classes/tables/scantable.py

@@ -1,4 +1,4 @@
-from typing import Any, List
+from typing import Any, List, Optional
 
 from dedocutils.data_structures import BBox
 
@@ -10,7 +10,8 @@
 
 
 class ScanTable:
-    def __init__(self, page_number: int, matrix_cells: List[List[Cell]] = None, bbox: BBox = None, name: str = "", order: int = -1) -> None:
+    def __init__(self, page_number: int, matrix_cells: Optional[List[List[CellWithMeta]]] = None, bbox: Optional[BBox] = None,
+                 name: str = "", order: int = -1) -> None:
         self.matrix_cells = matrix_cells
         self.page_number = page_number
         self.locations = []
@@ -27,6 +28,15 @@ def extended(self, table: "ScanTable") -> None:
         # extend order
         self.order = max(self.order, table.order)
 
+    def check_on_cell_instance(self) -> bool:
+        if len(self.matrix_cells) == 0:
+            return False
+        if len(self.matrix_cells[0]) == 0:
+            return False
+        if not isinstance(self.matrix_cells[0][0], Cell):
+            return False
+        return True
+
     def to_table(self) -> Table:
         metadata = TableMetadata(page_id=self.page_number, uid=self.name, rotated_angle=self.location.rotated_angle)
         cells_with_meta = [[CellWithMeta.create_from_cell(cell) for cell in row] for row in self.matrix_cells]

dedoc/readers/pdf_reader/pdf_image_reader/table_recognizer/table_extractors/concrete_extractors/multipage_table_extractor.py

@@ -1,3 +1,4 @@
+import copy
 import logging
 from typing import List
 
@@ -155,24 +156,26 @@ def __is_one_table(self, t1: ScanTable, t2: ScanTable) -> bool:
         # condition 2. Exclusion of the duplicated header (if any)
         attr1 = TableAttributeExtractor.get_header_table(t1.matrix_cells)
         attr2 = TableAttributeExtractor.get_header_table(t2.matrix_cells)
+        t2_update = copy.deepcopy(t2)
         if TableAttributeExtractor.is_equal_attributes(attr1, attr2):
-            t2.matrix_cells = t2.matrix_cells[len(attr2):]
+            t2_update.matrix_cells = t2_update.matrix_cells[len(attr2):]
 
-        if len(t2.matrix_cells) == 0 or len(t1.matrix_cells) == 0:
+        if len(t2_update.matrix_cells) == 0 or len(t1.matrix_cells) == 0:
             return False
 
-        TableAttributeExtractor.clear_attributes(t2.matrix_cells)
+        TableAttributeExtractor.clear_attributes(t2_update.matrix_cells)
 
         # condition 3. Number of columns should be equal
-        if len(t1.matrix_cells[-1]) != len(t2.matrix_cells[0]):
+        if len(t1.matrix_cells[-1]) != len(t2_update.matrix_cells[0]):
             if self.config.get("debug_mode", False):
                 self.logger.debug("Different count column")
             return False
 
         # condition 4. Comparison of the widths of last and first rows
-        if not self.__is_equal_width_cells(t1.matrix_cells, t2.matrix_cells):
+        if t1.check_on_cell_instance() and t2_update.check_on_cell_instance() and not self.__is_equal_width_cells(t1.matrix_cells, t2_update.matrix_cells):
             if self.config.get("debug_mode", False):
                 self.logger.debug("Different width columns")
             return False
 
+        t2.matrix_cells = copy.deepcopy(t2_update.matrix_cells)  # save changes
         return True

dedoc/readers/pdf_reader/pdf_image_reader/table_recognizer/table_extractors/concrete_extractors/table_attribute_extractor.py

@@ -31,8 +31,22 @@ def is_equal_attributes(attr1: List[List[Cell]], attr2: List[List[Cell]], thr_si
 
         return True
 
+    @staticmethod
+    def check_have_attributes(matrix_table: List[List[Cell]]) -> bool:
+        if len(matrix_table) == 0:
+            return False
+        if len(matrix_table[0]) == 0:
+            return False
+        if not hasattr(matrix_table[0][0], "is_attribute"):
+            return False
+        return True
+
     @staticmethod
     def get_header_table(matrix_table: List[List[Cell]]) -> List[List[Cell]]:
+
+        if not TableAttributeExtractor.check_have_attributes(matrix_table):
+            return matrix_table[:1]
+
         header_rows = len(matrix_table)
         for (i, row) in enumerate(matrix_table):
             attrs = [cell for cell in row if cell.is_attribute]
@@ -44,6 +58,9 @@ def get_header_table(matrix_table: List[List[Cell]]) -> List[List[Cell]]:
 
     @staticmethod
     def clear_attributes(matrix_table: List[List[Cell]]) -> None:
+        if not TableAttributeExtractor.check_have_attributes(matrix_table):
+            return
+
         for row in matrix_table:
             for cell in row:
                 cell.is_attribute = False

dedoc/readers/pdf_reader/pdf_txtlayer_reader/pdf_tabby_reader.py

@@ -63,13 +63,7 @@ def read(self, file_path: str, parameters: Optional[dict] = None) -> Unstructure
         warnings = []
 
         with tempfile.TemporaryDirectory() as tmp_dir:
-            lines, tables, tables_on_images, attachments, document_metadata = self.__extract(
-                path=file_path,
-                parameters=parameters,
-                warnings=warnings,
-                tmp_dir=tmp_dir
-            )
-        lines = self.linker.link_objects(lines=lines, tables=tables_on_images, images=attachments)
+            lines, tables, attachments, document_metadata = self.__extract(path=file_path, parameters=parameters, warnings=warnings, tmp_dir=tmp_dir)
 
         if get_param_with_attachments(parameters) and self.attachment_extractor.can_extract(file_path):
             attachments += self.attachment_extractor.extract(file_path=file_path, parameters=parameters)
@@ -80,15 +74,15 @@ def read(self, file_path: str, parameters: Optional[dict] = None) -> Unstructure
 
         return self._postprocess(result)
 
-    def __extract(self, path: str, parameters: dict, warnings: list, tmp_dir: str)\
-            -> Tuple[List[LineWithMeta], List[Table], List[ScanTable], List[PdfImageAttachment], Optional[dict]]:
+    def __extract(self, path: str, parameters: dict, warnings: List[str], tmp_dir: str)\
+            -> Tuple[List[LineWithMeta], List[Table], List[PdfImageAttachment], Optional[dict]]:
         import math
         from dedoc.utils.pdf_utils import get_pdf_page_count
         from dedoc.utils.utils import calculate_file_hash
         from dedoc.utils.parameter_utils import get_param_page_slice, get_param_with_attachments
         from dedoc.utils.parameter_utils import get_param_need_gost_frame_analysis
 
-        all_lines, all_tables, all_tables_on_images, all_attached_images = [], [], [], []
+        all_lines, all_tables, all_scan_tables, all_attached_images = [], [], [], []
         with_attachments = get_param_with_attachments(parameters)
         document_metadata = None
 
@@ -106,12 +100,11 @@ def __extract(self, path: str, parameters: dict, warnings: list, tmp_dir: str)\
                 document_metadata["last_page"] = last_page
 
             if empty_page_limit:
-                return all_lines, all_tables, all_tables_on_images, all_attached_images, document_metadata
+                return all_lines, all_tables, all_attached_images, document_metadata
 
-        gost_json_path = ""
-        remove_frame = get_param_need_gost_frame_analysis(parameters)
-        if remove_frame:
-            gost_json_path = self.__save_gost_frame_boxes_to_json(first_page=first_page, last_page=last_page, page_count=page_count, tmp_dir=tmp_dir, path=path)
+        remove_gost_frame = get_param_need_gost_frame_analysis(parameters)
+        gost_json_path = self.__save_gost_frame_boxes_to_json(first_page=first_page, last_page=last_page, page_count=page_count, tmp_dir=tmp_dir, path=path) \
+            if remove_gost_frame else ""
 
         # in java tabby reader page numeration starts with 1, end_page is included
         first_tabby_page = first_page + 1 if first_page is not None else 1
@@ -122,24 +115,26 @@ def __extract(self, path: str, parameters: dict, warnings: list, tmp_dir: str)\
                                       end_page=last_tabby_page,
                                       tmp_dir=tmp_dir,
                                       gost_json_path=gost_json_path,
-                                      remove_frame=remove_frame)
+                                      remove_frame=remove_gost_frame)
 
         pages = document.get("pages", [])
         for page in pages:
             page_lines = self.__get_lines_with_location(page, file_hash)
             if page_lines:
                 all_lines.extend(page_lines)
-            page_tables, table_on_images = self.__get_tables(page)
-            assert len(page_tables) == len(table_on_images)
-            if page_tables:
-                all_tables.extend(page_tables)
-                all_tables_on_images.extend(table_on_images)
+            scan_tables = self.__get_tables(page)
+            all_scan_tables.extend(scan_tables)
 
             attached_images = self.__get_attached_images(page=page, parameters=parameters, path=path) if with_attachments else []
             if attached_images:
                 all_attached_images.extend(attached_images)
 
-        return all_lines, all_tables, all_tables_on_images, all_attached_images, document_metadata
+        mp_tables = self.table_recognizer.convert_to_multipages_tables(all_scan_tables, lines_with_meta=all_lines)
+        all_lines = self.linker.link_objects(lines=all_lines, tables=mp_tables, images=all_attached_images)
+
+        tables = [scan_table.to_table() for scan_table in mp_tables]
+
+        return all_lines, tables, all_attached_images, document_metadata
 
     def __save_gost_frame_boxes_to_json(self, first_page: Optional[int], last_page: Optional[int], page_count: int, path: str, tmp_dir: str) -> str:
         from joblib import Parallel, delayed
@@ -150,24 +145,25 @@ def __save_gost_frame_boxes_to_json(self, first_page: Optional[int], last_page:
         images = self._get_images(path, first_page, last_page)
 
         gost_analyzed_images = Parallel(n_jobs=self.config["n_jobs"])(delayed(self.gost_frame_recognizer.rec_and_clean_frame)(image) for image in images)
+
         result_dict = {
             page_number: {**page_data[1].to_dict(), **{"original_image_width": page_data[2][1], "original_image_height": page_data[2][0]}}
             for page_number, page_data in enumerate(gost_analyzed_images, start=first_page)
         }
+
         result_json_path = os.path.join(tmp_dir, "gost_frame_bboxes.json")
         with open(result_json_path, "w") as f:
             json.dump(result_dict, f)
+
         return result_json_path
 
-    def __get_tables(self, page: dict) -> Tuple[List[Table], List[ScanTable]]:
+    def __get_tables(self, page: dict) -> List[ScanTable]:
         import uuid
         from dedoc.data_structures.cell_with_meta import CellWithMeta
         from dedoc.data_structures.concrete_annotations.bbox_annotation import BBoxAnnotation
         from dedoc.data_structures.line_metadata import LineMetadata
-        from dedoc.data_structures.table_metadata import TableMetadata
 
-        tables = []
-        tables_on_image = []
+        scan_tables = []
         page_number = page["number"]
         page_width = int(page["width"])
         page_height = int(page["height"])
@@ -179,7 +175,7 @@ def __get_tables(self, page: dict) -> Tuple[List[Table], List[ScanTable]]:
             cell_properties = table["cell_properties"]
             assert len(rows) == len(cell_properties)
 
-            result_cells = []
+            cells = []
             for num_row, row in enumerate(rows):
                 assert len(row) == len(cell_properties[num_row])
 
@@ -191,20 +187,22 @@ def __get_tables(self, page: dict) -> Tuple[List[Table], List[ScanTable]]:
                     for c in cell_blocks:
                         cell_bbox = BBox(x_top_left=int(c["x_top_left"]), y_top_left=int(c["y_top_left"]), width=int(c["width"]), height=int(c["height"]))
                         annotations.append(BBoxAnnotation(c["start"], c["end"], cell_bbox, page_width=page_width, page_height=page_height))
+                    """
+                        TODO: change to Cell class after tabby can return cell coordinates. Then set type Cell in class "ScanTable"
+                        https://jira.intra.ispras.ru/browse/TLDR-851
+                    """
 
                     result_row.append(CellWithMeta(
                         lines=[LineWithMeta(line=cell["text"], metadata=LineMetadata(page_id=page_number, line_id=0), annotations=annotations)],
                         colspan=cell_properties[num_row][num_col]["col_span"],
                         rowspan=cell_properties[num_row][num_col]["row_span"],
                         invisible=bool(cell_properties[num_row][num_col]["invisible"])
                     ))
-                result_cells.append(result_row)
+                cells.append(result_row)
 
-            table_name = str(uuid.uuid4())
-            tables.append(Table(cells=result_cells, metadata=TableMetadata(page_id=page_number, uid=table_name)))
-            tables_on_image.append(ScanTable(page_number=page_number, matrix_cells=None, bbox=table_bbox, name=table_name, order=order))
+            scan_tables.append(ScanTable(page_number=page_number, matrix_cells=cells, bbox=table_bbox, name=str(uuid.uuid4()), order=order))
 
-        return tables, tables_on_image
+        return scan_tables
 
     def __get_attached_images(self, page: dict, parameters: dict, path: str) -> List[PdfImageAttachment]:
         import os