RED-6009: Document Tree Structure

* refactored data mapper classes
* refactored graph construction to remove the possibility of adding anythingbut table cells underneath a table node

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/classification/parsing/PDFAreaTextStripper.java

@@ -1,82 +0,0 @@
-package com.iqser.red.service.redaction.v1.server.layoutparsing.classification.parsing;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-
-import org.apache.pdfbox.text.PDFTextStripperByArea;
-import org.apache.pdfbox.text.TextPosition;
-
-import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.model.text.TextPositionSequence;
-
-import lombok.Getter;
-import lombok.Setter;
-
-public class PDFAreaTextStripper extends PDFTextStripperByArea {
-
-    @Getter
-    private List<TextPositionSequence> textPositionSequences = new ArrayList<>();
-
-    @Setter
-    private int pageNumber;
-
-
-    public PDFAreaTextStripper() throws IOException {
-
-    }
-
-
-    @Override
-    public void writeString(String text, List<TextPosition> textPositions) throws IOException {
-
-        int startIndex = 0;
-        for (int i = 0; i <= textPositions.size() - 1; i++) {
-
-            if (i == 0 && (textPositions.get(i).getUnicode().equals(" ") || textPositions.get(i).getUnicode().equals("\u00A0"))) {
-                startIndex++;
-                continue;
-            }
-
-            // Strange but sometimes this is happening, for example: Metolachlor2.pdf
-            if (i > 0 && textPositions.get(i).getX() < textPositions.get(i - 1).getX()) {
-                List<TextPosition> sublist = textPositions.subList(startIndex, i);
-                if (!(sublist.isEmpty() || sublist.size() == 1 && (sublist.get(0).getUnicode().equals(" ") || sublist.get(0).getUnicode().equals("\u00A0")))) {
-                    textPositionSequences.add(new TextPositionSequence(sublist, pageNumber));
-                }
-                startIndex = i;
-            }
-
-            if (textPositions.get(i).getRotation() == 0 && i > 0 && textPositions.get(i).getX() > textPositions.get(i - 1).getEndX() + 1) {
-                List<TextPosition> sublist = textPositions.subList(startIndex, i);
-                if (!(sublist.isEmpty() || sublist.size() == 1 && (sublist.get(0).getUnicode().equals(" ") || sublist.get(0).getUnicode().equals("\u00A0")))) {
-                    textPositionSequences.add(new TextPositionSequence(sublist, pageNumber));
-                }
-                startIndex = i;
-            }
-
-            if (i > 0 && (textPositions.get(i).getUnicode().equals(" ") || textPositions.get(i).getUnicode().equals("\u00A0")) && i <= textPositions.size() - 2) {
-                List<TextPosition> sublist = textPositions.subList(startIndex, i);
-                if (!(sublist.isEmpty() || sublist.size() == 1 && (sublist.get(0).getUnicode().equals(" ") || sublist.get(0).getUnicode().equals("\u00A0")))) {
-                    textPositionSequences.add(new TextPositionSequence(sublist, pageNumber));
-                }
-                startIndex = i + 1;
-            }
-        }
-
-        List<TextPosition> sublist = textPositions.subList(startIndex, textPositions.size());
-        if (!sublist.isEmpty() && (sublist.get(sublist.size() - 1).getUnicode().equals(" ") || sublist.get(sublist.size() - 1).getUnicode().equals("\u00A0"))) {
-            sublist = sublist.subList(0, sublist.size() - 1);
-        }
-        if (!(sublist.isEmpty() || sublist.size() == 1 && (sublist.get(0).getUnicode().equals(" ") || sublist.get(0).getUnicode().equals("\u00A0")))) {
-            textPositionSequences.add(new TextPositionSequence(sublist, pageNumber));
-        }
-        super.writeString(text);
-    }
-
-
-    public void clearPositions() {
-
-        textPositionSequences = new ArrayList<>();
-    }
-
-}

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/data/AtomicPositionBlockData.java

@@ -1,5 +1,10 @@
 package com.iqser.red.service.redaction.v1.server.layoutparsing.document.data;
 
+import java.awt.geom.Rectangle2D;
+import java.util.List;
+
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.textblock.AtomicTextBlock;
+
 import lombok.AccessLevel;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
@@ -16,4 +21,40 @@ public class AtomicPositionBlockData {
     int[] stringIdxToPositionIdx;
     float[][] positions;
 
+
+    public static AtomicPositionBlockData fromAtomicTextBlock(AtomicTextBlock atomicTextBlock) {
+
+        return AtomicPositionBlockData.builder()
+                .id(atomicTextBlock.getId())
+                .positions(toPrimitiveFloatMatrix(atomicTextBlock.getPositions()))
+                .stringIdxToPositionIdx(toPrimitiveIntArray(atomicTextBlock.getStringIdxToPositionIdx()))
+                .build();
+    }
+
+
+    private static int[] toPrimitiveIntArray(List<Integer> list) {
+
+        int[] array = new int[list.size()];
+        for (int i = 0; i < list.size(); i++) {
+            array[i] = list.get(i);
+        }
+        return array;
+    }
+
+
+    private static float[][] toPrimitiveFloatMatrix(List<Rectangle2D> positions) {
+
+        float[][] positionMatrix = new float[positions.size()][];
+        for (int i = 0; i < positions.size(); i++) {
+            positionMatrix[i] = toArray(positions.get(i));
+        }
+        return positionMatrix;
+    }
+
+
+    private static float[] toArray(Rectangle2D positions) {
+
+        return new float[]{(float) positions.getMinX(), (float) positions.getMinY(), (float) positions.getWidth(), (float) positions.getHeight()};
+    }
+
 }

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/data/AtomicTextBlockData.java

@@ -1,5 +1,9 @@
 package com.iqser.red.service.redaction.v1.server.layoutparsing.document.data;
 
+import java.util.List;
+
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.textblock.AtomicTextBlock;
+
 import lombok.AccessLevel;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
@@ -20,4 +24,28 @@ public class AtomicTextBlockData {
     int end;
     int[] lineBreaks;
 
+
+    public static AtomicTextBlockData fromAtomicTextBlock(AtomicTextBlock atomicTextBlock) {
+
+        return AtomicTextBlockData.builder()
+                .id(atomicTextBlock.getId())
+                .page(atomicTextBlock.getPage().getNumber().longValue())
+                .searchText(atomicTextBlock.getSearchText())
+                .numberOnPage(atomicTextBlock.getNumberOnPage())
+                .start(atomicTextBlock.getBoundary().start())
+                .end(atomicTextBlock.getBoundary().end())
+                .lineBreaks(toPrimitiveIntArray(atomicTextBlock.getLineBreaks()))
+                .build();
+    }
+
+
+    private static int[] toPrimitiveIntArray(List<Integer> list) {
+
+        int[] array = new int[list.size()];
+        for (int i = 0; i < list.size(); i++) {
+            array[i] = list.get(i);
+        }
+        return array;
+    }
+
 }

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/data/mapper/DocumentDataMapper.java

@@ -1,6 +1,5 @@
 package com.iqser.red.service.redaction.v1.server.layoutparsing.document.data.mapper;
 
-import java.awt.geom.Rectangle2D;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -30,13 +29,13 @@ public class DocumentDataMapper {
         List<AtomicTextBlockData> atomicTextBlockData = document.streamTerminalTextBlocksInOrder()
                 .flatMap(textBlock -> textBlock.getAtomicTextBlocks().stream())
                 .distinct()
-                .map(DocumentDataMapper::toAtomicTextBlockData)
+                .map(AtomicTextBlockData::fromAtomicTextBlock)
                 .toList();
 
         List<AtomicPositionBlockData> atomicPositionBlockData = document.streamTerminalTextBlocksInOrder()
                 .flatMap(textBlock -> textBlock.getAtomicTextBlocks().stream())
                 .distinct()
-                .map(DocumentDataMapper::toAtomicPositionBlockData)
+                .map(AtomicPositionBlockData::fromAtomicTextBlock)
                 .toList();
 
         List<PageData> pageData = document.getPages().stream().map(DocumentDataMapper::toPageData).toList();
@@ -75,7 +74,7 @@ public class DocumentDataMapper {
         };
 
         return DocumentTreeData.EntryData.builder()
-                .treeId(toPrimitiveIntArray(entry.getTocId()))
+                .treeId(toPrimitiveIntArray(entry.getTreeId()))
                 .subEntries(entry.getChildren().stream().map(DocumentDataMapper::toEntryData).toList())
                 .type(entry.getType())
                 .atomicBlocks(atomicTextBlocks)
@@ -97,46 +96,7 @@ public class DocumentDataMapper {
     }
 
 
-    private AtomicTextBlockData toAtomicTextBlockData(AtomicTextBlock atomicTextBlock) {
+    private static int[] toPrimitiveIntArray(List<Integer> list) {
-
-        return AtomicTextBlockData.builder()
-                .id(atomicTextBlock.getId())
-                .page(atomicTextBlock.getPage().getNumber().longValue())
-                .searchText(atomicTextBlock.getSearchText())
-                .numberOnPage(atomicTextBlock.getNumberOnPage())
-                .start(atomicTextBlock.getBoundary().start())
-                .end(atomicTextBlock.getBoundary().end())
-                .lineBreaks(toPrimitiveIntArray(atomicTextBlock.getLineBreaks()))
-                .build();
-    }
-
-
-    private AtomicPositionBlockData toAtomicPositionBlockData(AtomicTextBlock atomicTextBlock) {
-
-        return AtomicPositionBlockData.builder()
-                .id(atomicTextBlock.getId())
-                .positions(toPrimitiveFloatMatrix(atomicTextBlock.getPositions()))
-                .stringIdxToPositionIdx(toPrimitiveIntArray(atomicTextBlock.getStringIdxToPositionIdx()))
-                .build();
-    }
-
-
-    private float[][] toPrimitiveFloatMatrix(List<Rectangle2D> positions) {
-
-        float[][] positionMatrix = new float[positions.size()][];
-        for (int i = 0; i < positions.size(); i++) {
-            float[] singlePositions = new float[4];
-            singlePositions[0] = (float) positions.get(i).getMinX();
-            singlePositions[1] = (float) positions.get(i).getMinY();
-            singlePositions[2] = (float) positions.get(i).getWidth();
-            singlePositions[3] = (float) positions.get(i).getHeight();
-            positionMatrix[i] = singlePositions;
-        }
-        return positionMatrix;
-    }
-
-
-    private int[] toPrimitiveIntArray(List<Integer> list) {
 
         int[] array = new int[list.size()];
         for (int i = 0; i < list.size(); i++) {

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/data/mapper/DocumentGraphMapper.java

@@ -3,7 +3,6 @@ package com.iqser.red.service.redaction.v1.server.layoutparsing.document.data.ma
 import static com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.NodeType.FOOTER;
 import static com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.NodeType.HEADER;
 
-import java.awt.geom.Rectangle2D;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.LinkedList;
@@ -11,13 +10,11 @@ import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
 
-import com.google.common.primitives.Ints;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.data.AtomicPositionBlockData;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.data.AtomicTextBlockData;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.data.DocumentData;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.data.DocumentTreeData;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.data.PageData;
-import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.Boundary;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.DocumentTree;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Document;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Footer;
@@ -86,8 +83,8 @@ public class DocumentGraphMapper {
                 TextBlock textBlock = toTextBlock(entryData.getAtomicBlocks(), context, node);
                 node.setLeafTextBlock(textBlock);
             }
-            List<Integer> tocId = Arrays.stream(entryData.getTreeId()).boxed().toList();
+            List<Integer> treeId = Arrays.stream(entryData.getTreeId()).boxed().toList();
-            node.setTreeId(tocId);
+            node.setTreeId(treeId);
 
             if (entryData.getType() == HEADER) {
                 pages.forEach(page -> page.setHeader((Header) node));
@@ -96,7 +93,7 @@ public class DocumentGraphMapper {
             } else {
                 pages.forEach(page -> page.getMainBody().add(node));
             }
-            newEntries.add(DocumentTree.Entry.builder().tocId(tocId).type(entryData.getType()).children(buildEntries(entryData.getSubEntries(), context)).node(node).build());
+            newEntries.add(DocumentTree.Entry.builder().treeId(treeId).type(entryData.getType()).children(buildEntries(entryData.getSubEntries(), context)).node(node).build());
         }
         return newEntries;
     }
@@ -168,12 +165,16 @@ public class DocumentGraphMapper {
 
     private TextBlock toTextBlock(Long[] atomicTextBlockIds, Context context, SemanticNode parent) {
 
-        return Arrays.stream(atomicTextBlockIds)
+        return Arrays.stream(atomicTextBlockIds).map(atomicTextBlockId -> getAtomicTextBlock(context, parent, atomicTextBlockId)).collect(new TextBlockCollector());
-                .map(atomicTextBlockId -> toAtomicTextBlock(context.atomicTextBlockData.get(Math.toIntExact(atomicTextBlockId)),
+    }
-                        context.atomicPositionBlockData.get(Math.toIntExact(atomicTextBlockId)),
+
-                        parent,
+
-                        context))
+    private static AtomicTextBlock getAtomicTextBlock(Context context, SemanticNode parent, Long atomicTextBlockId) {
-                .collect(new TextBlockCollector());
+
+        return AtomicTextBlock.fromAtomicTextBlockData(context.atomicTextBlockData.get(Math.toIntExact(atomicTextBlockId)),
+                context.atomicPositionBlockData.get(Math.toIntExact(atomicTextBlockId)),
+                parent,
+                getPage(context.atomicTextBlockData.get(Math.toIntExact(atomicTextBlockId)).getPage(), context));
     }
 
 
@@ -183,28 +184,6 @@ public class DocumentGraphMapper {
     }
 
 
-    private AtomicTextBlock toAtomicTextBlock(AtomicTextBlockData atomicTextBlockData, AtomicPositionBlockData atomicPositionBlockData, SemanticNode parent, Context context) {
-
-        return AtomicTextBlock.builder()
-                .id(atomicTextBlockData.getId())
-                .numberOnPage(atomicTextBlockData.getNumberOnPage())
-                .page(getPage(atomicTextBlockData.getPage(), context))
-                .boundary(new Boundary(atomicTextBlockData.getStart(), atomicTextBlockData.getEnd()))
-                .searchText(atomicTextBlockData.getSearchText())
-                .lineBreaks(Ints.asList(atomicTextBlockData.getLineBreaks()))
-                .stringIdxToPositionIdx(Ints.asList(atomicPositionBlockData.getStringIdxToPositionIdx()))
-                .positions(toRectangle2DList(atomicPositionBlockData.getPositions()))
-                .parent(parent)
-                .build();
-    }
-
-
-    private static List<Rectangle2D> toRectangle2DList(float[][] positions) {
-
-        return Arrays.stream(positions).map(floatArr -> (Rectangle2D) new Rectangle2D.Float(floatArr[0], floatArr[1], floatArr[2], floatArr[3])).toList();
-    }
-
-
     private Page getPage(Long pageIndex, Context context) {
 
         return context.pages.stream()

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/factory/DocumentGraphFactory.java

@@ -23,6 +23,7 @@ import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.mo
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.DocumentTree;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Document;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Footer;
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.GenericSemanticNode;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Header;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Headline;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Image;
@@ -30,7 +31,6 @@ import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.no
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Page;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Paragraph;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Section;
-import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.SemanticNode;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.textblock.AtomicTextBlock;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.utils.TextPositionOperations;
 import com.iqser.red.service.redaction.v1.server.redaction.utils.IdBuilder;
@@ -75,11 +75,14 @@ public class DocumentGraphFactory {
     }
 
 
-    public void addParagraphOrHeadline(SemanticNode parentNode, ClassificationTextBlock originalTextBlock, Context context, List<ClassificationTextBlock> textBlocksToMerge) {
+    public void addParagraphOrHeadline(GenericSemanticNode parentNode,
+                                       ClassificationTextBlock originalTextBlock,
+                                       Context context,
+                                       List<ClassificationTextBlock> textBlocksToMerge) {
 
         Page page = context.getPage(originalTextBlock.getPage());
 
-        SemanticNode node;
+        GenericSemanticNode node;
         if (pageBlockIsHeadline(originalTextBlock)) {
             node = Headline.builder().documentTree(context.getDocumentTree()).build();
         } else {
@@ -93,12 +96,12 @@ public class DocumentGraphFactory {
         AtomicTextBlock textBlock = context.textBlockFactory.buildAtomicTextBlock(TextPositionOperations.mergeAndSortTextPositionSequenceByYThenX(textBlocks), node, context, page);
 
         if (node instanceof Headline headline) {
-            List<Integer> tocId = context.documentTree.createNewChildEntryAndReturnId(parentNode.getTreeId(), NodeType.HEADLINE, node);
+            List<Integer> tocId = context.documentTree.createNewChildEntryAndReturnId(parentNode, NodeType.HEADLINE, node);
             headline.setLeafTextBlock(textBlock);
             headline.setTreeId(tocId);
         }
         if (node instanceof Paragraph paragraph) {
-            List<Integer> tocId = context.documentTree.createNewChildEntryAndReturnId(parentNode.getTreeId(), NodeType.PARAGRAPH, node);
+            List<Integer> tocId = context.documentTree.createNewChildEntryAndReturnId(parentNode, NodeType.PARAGRAPH, node);
             paragraph.setLeafTextBlock(textBlock);
             paragraph.setTreeId(tocId);
         }
@@ -119,7 +122,7 @@ public class DocumentGraphFactory {
                 .build();
         page.getMainBody().add(imageNode);
 
-        List<Integer> tocId = context.getDocumentTree().createNewChildEntryAndReturnId(section.getTreeId(), NodeType.IMAGE, imageNode);
+        List<Integer> tocId = context.getDocumentTree().createNewChildEntryAndReturnId(section, NodeType.IMAGE, imageNode);
         imageNode.setTreeId(tocId);
     }
 

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/factory/SectionNodeFactory.java

@@ -14,10 +14,10 @@ import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.mo
 import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.model.image.ClassifiedImage;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.model.table.TablePageBlock;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.model.text.ClassificationTextBlock;
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.GenericSemanticNode;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.NodeType;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Page;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Section;
-import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.SemanticNode;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.utils.TableMergingUtility;
 
 import lombok.experimental.UtilityClass;
@@ -25,7 +25,7 @@ import lombok.experimental.UtilityClass;
 @UtilityClass
 public class SectionNodeFactory {
 
-    public void addSection(SemanticNode parentNode, List<AbstractPageBlock> pageBlocks, List<ClassifiedImage> images, DocumentGraphFactory.Context context) {
+    public void addSection(GenericSemanticNode parentNode, List<AbstractPageBlock> pageBlocks, List<ClassifiedImage> images, DocumentGraphFactory.Context context) {
 
         if (pageBlocks.isEmpty()) {
             return;
@@ -40,7 +40,7 @@ public class SectionNodeFactory {
         if (parentNode == null) {
             tocId = context.getDocumentTree().createNewMainEntryAndReturnId(NodeType.SECTION, section);
         } else {
-            tocId = context.getDocumentTree().createNewChildEntryAndReturnId(parentNode.getTreeId(), NodeType.SECTION, section);
+            tocId = context.getDocumentTree().createNewChildEntryAndReturnId(parentNode, NodeType.SECTION, section);
         }
         section.setTreeId(tocId);
 

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/factory/TableNodeFactory.java

@@ -11,11 +11,13 @@ import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.mo
 import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.model.table.Cell;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.model.table.TablePageBlock;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.model.text.TextPositionSequence;
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.GenericSemanticNode;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.NodeType;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Page;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.SemanticNode;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Table;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.TableCell;
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.textblock.TextBlock;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.utils.TextPositionOperations;
 
 import lombok.experimental.UtilityClass;
@@ -26,7 +28,7 @@ public class TableNodeFactory {
     public static final double TABLE_CELL_MERGE_CONTENTS_SIZE_THRESHOLD = 0.05;
 
 
-    public void addTable(SemanticNode parentNode, List<TablePageBlock> tablesToMerge, DocumentGraphFactory.Context context) {
+    public void addTable(GenericSemanticNode parentNode, List<TablePageBlock> tablesToMerge, DocumentGraphFactory.Context context) {
 
         setPageNumberInCellTextBlocks(tablesToMerge);
         Set<Page> pages = tablesToMerge.stream().map(AbstractPageBlock::getPage).map(context::getPage).collect(Collectors.toSet());
@@ -35,7 +37,7 @@ public class TableNodeFactory {
 
         pages.forEach(page -> addTableToPage(page, parentNode, table));
 
-        List<Integer> tocId = context.getDocumentTree().createNewChildEntryAndReturnId(parentNode.getTreeId(), NodeType.TABLE, table);
+        List<Integer> tocId = context.getDocumentTree().createNewChildEntryAndReturnId(parentNode, NodeType.TABLE, table);
         table.setTreeId(tocId);
         addTableCells(mergedRows, table, context);
 
@@ -82,20 +84,20 @@ public class TableNodeFactory {
 
 
     @SuppressWarnings("PMD.UnusedPrivateMethod") // PMD actually flags this wrong
-    private void addTableCell(Cell cell, int rowIndex, int colIndex, SemanticNode parentNode, DocumentGraphFactory.Context context) {
+    private void addTableCell(Cell cell, int rowIndex, int colIndex, Table tableNode, DocumentGraphFactory.Context context) {
 
         Page page = context.getPage(cell.getPageNumber());
 
         TableCell tableCell = TableCell.builder().documentTree(context.getDocumentTree()).row(rowIndex).col(colIndex).header(cell.isHeaderCell()).bBox(cell.getBounds2D()).build();
         page.getMainBody().add(tableCell);
 
-        com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.textblock.TextBlock textBlock;
+        TextBlock textBlock;
 
-        List<Integer> tocId = context.getDocumentTree().createNewChildEntryAndReturnId(parentNode.getTreeId(), NodeType.TABLE_CELL, tableCell);
+        List<Integer> tocId = context.getDocumentTree().createNewTableChildEntryAndReturnId(tableNode, NodeType.TABLE_CELL, tableCell);
         tableCell.setTreeId(tocId);
 
         if (cell.getTextBlocks().isEmpty()) {
-            tableCell.setLeafTextBlock(context.getTextBlockFactory().emptyTextBlock(parentNode, context, page));
+            tableCell.setLeafTextBlock(context.getTextBlockFactory().emptyTextBlock(tableNode, context, page));
             tableCell.setLeaf(true);
 
         } else if (cell.getTextBlocks().size() == 1) {

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/DocumentTree.java

@@ -10,8 +10,11 @@ import java.util.stream.Stream;
 
 import com.google.common.hash.Hashing;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Document;
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.GenericSemanticNode;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.NodeType;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.SemanticNode;
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Table;
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.TableCell;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.textblock.TextBlock;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.textblock.TextBlockCollector;
 
@@ -30,7 +33,7 @@ public class DocumentTree {
 
     public DocumentTree(Document document) {
 
-        root = Entry.builder().tocId(Collections.emptyList()).type(NodeType.DOCUMENT).children(new LinkedList<>()).node(document).build();
+        root = Entry.builder().treeId(Collections.emptyList()).type(NodeType.DOCUMENT).children(new LinkedList<>()).node(document).build();
     }
 
 
@@ -40,13 +43,31 @@ public class DocumentTree {
     }
 
 
-    public List<Integer> createNewMainEntryAndReturnId(NodeType nodeType, SemanticNode node) {
+    public List<Integer> createNewMainEntryAndReturnId(NodeType nodeType, GenericSemanticNode node) {
 
-        return createNewChildEntryAndReturnId(Collections.emptyList(), nodeType, node);
+        return createNewChildEntryAndReturnIdImpl(Collections.emptyList(), nodeType, node);
     }
 
 
-    public List<Integer> createNewChildEntryAndReturnId(List<Integer> parentId, NodeType nodeType, SemanticNode node) {
+    public List<Integer> createNewChildEntryAndReturnId(GenericSemanticNode parentNode, NodeType nodeType, GenericSemanticNode node) {
+
+        return createNewChildEntryAndReturnIdImpl(parentNode.getTreeId(), nodeType, node);
+    }
+
+
+    public List<Integer> createNewChildEntryAndReturnId(GenericSemanticNode parentNode, NodeType nodeType, Table node) {
+
+        return createNewChildEntryAndReturnIdImpl(parentNode.getTreeId(), nodeType, node);
+    }
+
+
+    public List<Integer> createNewTableChildEntryAndReturnId(Table parentTable, NodeType nodeType, TableCell tableCell) {
+
+        return createNewChildEntryAndReturnIdImpl(parentTable.getTreeId(), nodeType, tableCell);
+    }
+
+
+    private List<Integer> createNewChildEntryAndReturnIdImpl(List<Integer> parentId, NodeType nodeType, SemanticNode node) {
 
         if (!entryExists(parentId)) {
             throw new UnsupportedOperationException(format("parentId %s does not exist!", parentId));
@@ -55,19 +76,19 @@ public class DocumentTree {
         Entry parent = getEntryById(parentId);
         List<Integer> newId = new LinkedList<>(parentId);
         newId.add(parent.children.size());
-        parent.children.add(Entry.builder().tocId(newId).node(node).type(nodeType).children(new LinkedList<>()).build());
+        parent.children.add(Entry.builder().treeId(newId).node(node).type(nodeType).children(new LinkedList<>()).build());
 
         return newId;
     }
 
 
-    private boolean entryExists(List<Integer> tocId) {
+    private boolean entryExists(List<Integer> treeId) {
 
-        if (tocId.isEmpty()) {
+        if (treeId.isEmpty()) {
             return root != null;
         }
-        Entry entry = root.children.get(tocId.get(0));
+        Entry entry = root.children.get(treeId.get(0));
-        for (int id : tocId.subList(1, tocId.size())) {
+        for (int id : treeId.subList(1, treeId.size())) {
             if (id >= entry.children.size() || 0 > id) {
                 return false;
             }
@@ -77,43 +98,49 @@ public class DocumentTree {
     }
 
 
-    public Entry getParentEntryById(List<Integer> tocId) {
+    public Entry getParentEntryById(List<Integer> treeId) {
 
-        return getEntryById(getParentId(tocId));
+        return getEntryById(getParentId(treeId));
     }
 
 
-    public boolean hasParentById(List<Integer> tocId) {
+    public boolean hasParentById(List<Integer> treeId) {
 
-        return !tocId.isEmpty();
+        return !treeId.isEmpty();
     }
 
 
-    public Stream<SemanticNode> streamChildrenNodes(List<Integer> tocId) {
+    public Stream<SemanticNode> streamChildNodes(List<Integer> treeId) {
 
-        return getEntryById(tocId).children.stream().map(Entry::getNode);
+        return getEntryById(treeId).children.stream().map(Entry::getNode);
     }
 
 
-    private static List<Integer> getParentId(List<Integer> tocId) {
+    public Stream<SemanticNode> streamChildNodesOfType(List<Integer> treeId, NodeType nodeType) {
 
-        if (tocId.isEmpty()) {
+        return getEntryById(treeId).children.stream().filter(entry -> entry.type.equals(nodeType)).map(Entry::getNode);
+    }
+
+
+    private static List<Integer> getParentId(List<Integer> treeId) {
+
+        if (treeId.isEmpty()) {
             throw new UnsupportedOperationException("Root has no parent!");
         }
-        if (tocId.size() < 2) {
+        if (treeId.size() < 2) {
             return Collections.emptyList();
         }
-        return tocId.subList(0, tocId.size() - 1);
+        return treeId.subList(0, treeId.size() - 1);
     }
 
 
-    public Entry getEntryById(List<Integer> tocId) {
+    public Entry getEntryById(List<Integer> treeId) {
 
-        if (tocId.isEmpty()) {
+        if (treeId.isEmpty()) {
             return root;
         }
-        Entry entry = root.children.get(tocId.get(0));
+        Entry entry = root.children.get(treeId.get(0));
-        for (int id : tocId.subList(1, tocId.size())) {
+        for (int id : treeId.subList(1, treeId.size())) {
             entry = entry.children.get(id);
         }
         return entry;
@@ -151,12 +178,12 @@ public class DocumentTree {
     }
 
 
-    public SemanticNode getHighestParentById(List<Integer> tocId) {
+    public SemanticNode getHighestParentById(List<Integer> treeId) {
 
-        if (tocId.isEmpty()) {
+        if (treeId.isEmpty()) {
             return root.node;
         }
-        return root.children.get(tocId.get(0)).node;
+        return root.children.get(treeId.get(0)).node;
     }
 
 
@@ -166,7 +193,7 @@ public class DocumentTree {
     @FieldDefaults(level = AccessLevel.PRIVATE, makeFinal = true)
     public static class Entry {
 
-        List<Integer> tocId;
+        List<Integer> treeId;
         NodeType type;
         SemanticNode node;
         List<Entry> children;

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/Document.java

@@ -28,7 +28,7 @@ import lombok.experimental.FieldDefaults;
 @AllArgsConstructor
 @NoArgsConstructor
 @FieldDefaults(level = AccessLevel.PRIVATE)
-public class Document implements SemanticNode {
+public class Document implements GenericSemanticNode {
 
     Set<Page> pages;
     DocumentTree documentTree;
@@ -46,7 +46,7 @@ public class Document implements SemanticNode {
 
     public List<Section> getMainSections() {
 
-        return streamChildren().filter(node -> node instanceof Section).map(node -> (Section) node).collect(Collectors.toList());
+        return streamChildrenOfType(NodeType.SECTION).map(node -> (Section) node).collect(Collectors.toList());
     }
 
 
@@ -73,10 +73,7 @@ public class Document implements SemanticNode {
     @Override
     public Headline getHeadline() {
 
-        return streamAllSubNodes().filter(node -> node instanceof Headline)
+        return streamChildrenOfType(NodeType.HEADLINE).map(node -> (Headline) node).findFirst().orElseThrow(() -> new NotFoundException("No Headlines found in this document!"));
-                .map(node -> (Headline) node)
-                .findFirst()
-                .orElseThrow(() -> new NotFoundException("No Headlines found in this document!"));
     }
 
 

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/Footer.java

@@ -21,7 +21,7 @@ import lombok.experimental.FieldDefaults;
 @AllArgsConstructor
 @NoArgsConstructor
 @FieldDefaults(level = AccessLevel.PRIVATE)
-public class Footer implements SemanticNode {
+public class Footer implements GenericSemanticNode {
 
     List<Integer> treeId;
     TextBlock leafTextBlock;

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/GenericSemanticNode.java

@@ -0,0 +1,5 @@
+package com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes;
+
+public interface GenericSemanticNode extends SemanticNode {
+
+}

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/Header.java

@@ -21,7 +21,7 @@ import lombok.experimental.FieldDefaults;
 @AllArgsConstructor
 @NoArgsConstructor
 @FieldDefaults(level = AccessLevel.PRIVATE)
-public class Header implements SemanticNode {
+public class Header implements GenericSemanticNode {
 
     List<Integer> treeId;
     TextBlock leafTextBlock;

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/Headline.java

@@ -21,13 +21,13 @@ import lombok.experimental.FieldDefaults;
 @AllArgsConstructor
 @NoArgsConstructor
 @FieldDefaults(level = AccessLevel.PRIVATE)
-public class Headline implements SemanticNode {
+public class Headline implements GenericSemanticNode {
 
     List<Integer> treeId;
     TextBlock leafTextBlock;
 
     @Builder.Default
-    boolean leaf = true;
+    final boolean leaf = true;
 
     @EqualsAndHashCode.Exclude
     DocumentTree documentTree;

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/Image.java

@@ -26,7 +26,7 @@ import lombok.experimental.FieldDefaults;
 @AllArgsConstructor
 @NoArgsConstructor
 @FieldDefaults(level = AccessLevel.PRIVATE)
-public class Image implements SemanticNode {
+public class Image implements GenericSemanticNode {
 
     List<Integer> treeId;
     String id;

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/Paragraph.java

@@ -19,7 +19,7 @@ import lombok.experimental.FieldDefaults;
 @Builder
 @AllArgsConstructor
 @FieldDefaults(level = AccessLevel.PRIVATE)
-public class Paragraph implements SemanticNode {
+public class Paragraph implements GenericSemanticNode {
 
     List<Integer> treeId;
     TextBlock leafTextBlock;

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/Section.java

@@ -22,7 +22,7 @@ import lombok.extern.slf4j.Slf4j;
 @Builder
 @AllArgsConstructor
 @FieldDefaults(level = AccessLevel.PRIVATE)
-public class Section implements SemanticNode {
+public class Section implements GenericSemanticNode {
 
     List<Integer> treeId;
 

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/SemanticNode.java

@@ -57,7 +57,7 @@ public interface SemanticNode {
     default Set<Page> getPages(Boundary boundary) {
 
         if (!getBoundary().contains(boundary)) {
-            throw new IllegalArgumentException(format("%s which was used to query foir pages is not contained in the %s of this node!", getBoundary(), boundary));
+            throw new IllegalArgumentException(format("%s which was used to query for pages is not contained in the %s of this node!", boundary, getBoundary()));
         }
         return buildTextBlock().getPages(boundary);
     }
@@ -211,7 +211,7 @@ public interface SemanticNode {
     default Integer getNumberOnPage() {
 
         TextBlock textBlock = buildTextBlock();
-        if (textBlock.getAtomicTextBlocks().size() > 0) {
+        if (!textBlock.getAtomicTextBlocks().isEmpty()) {
             return buildTextBlock().getAtomicTextBlocks().get(0).getNumberOnPage();
         } else {
             return -1;
@@ -226,7 +226,7 @@ public interface SemanticNode {
      */
     default boolean hasText() {
 
-        return buildTextBlock().length() > 0;
+        return !buildTextBlock().isEmpty();
     }
 
 
@@ -318,7 +318,18 @@ public interface SemanticNode {
      */
     default Stream<SemanticNode> streamChildren() {
 
-        return getDocumentTree().streamChildrenNodes(getTreeId());
+        return getDocumentTree().streamChildNodes(getTreeId());
+    }
+
+
+    /**
+     * Streams all children located directly underneath this node in the DocumentTree.
+     *
+     * @return Stream of all children
+     */
+    default Stream<SemanticNode> streamChildrenOfType(NodeType nodeType) {
+
+        return getDocumentTree().streamChildNodesOfType(getTreeId(), nodeType);
     }
 
 

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/Table.java

@@ -147,7 +147,7 @@ public class Table implements SemanticNode {
      */
     public Stream<TableCell> streamTableCells() {
 
-        return streamChildren().map(node -> (TableCell) node);
+        return streamChildrenOfType(NodeType.TABLE_CELL).map(node -> (TableCell) node);
     }
 
 

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/nodes/TableCell.java

@@ -23,7 +23,7 @@ import lombok.experimental.FieldDefaults;
 @Builder
 @AllArgsConstructor
 @FieldDefaults(level = AccessLevel.PRIVATE)
-public class TableCell implements SemanticNode {
+public class TableCell implements GenericSemanticNode {
 
     List<Integer> treeId;
     int row;

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/graph/textblock/AtomicTextBlock.java

@@ -3,12 +3,16 @@ package com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.t
 import static java.lang.String.format;
 
 import java.awt.geom.Rectangle2D;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import com.google.common.primitives.Ints;
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.data.AtomicPositionBlockData;
+import com.iqser.red.service.redaction.v1.server.layoutparsing.document.data.AtomicTextBlockData;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.Boundary;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.Page;
 import com.iqser.red.service.redaction.v1.server.layoutparsing.document.graph.nodes.SemanticNode;
@@ -51,6 +55,31 @@ public class AtomicTextBlock implements TextBlock {
     }
 
 
+    public static AtomicTextBlock fromAtomicTextBlockData(AtomicTextBlockData atomicTextBlockData,
+                                                          AtomicPositionBlockData atomicPositionBlockData,
+                                                          SemanticNode parent,
+                                                          Page page) {
+
+        return AtomicTextBlock.builder()
+                .id(atomicTextBlockData.getId())
+                .numberOnPage(atomicTextBlockData.getNumberOnPage())
+                .page(page)
+                .boundary(new Boundary(atomicTextBlockData.getStart(), atomicTextBlockData.getEnd()))
+                .searchText(atomicTextBlockData.getSearchText())
+                .lineBreaks(Ints.asList(atomicTextBlockData.getLineBreaks()))
+                .stringIdxToPositionIdx(Ints.asList(atomicPositionBlockData.getStringIdxToPositionIdx()))
+                .positions(toRectangle2DList(atomicPositionBlockData.getPositions()))
+                .parent(parent)
+                .build();
+    }
+
+
+    private static List<Rectangle2D> toRectangle2DList(float[][] positions) {
+
+        return Arrays.stream(positions).map(floatArr -> (Rectangle2D) new Rectangle2D.Float(floatArr[0], floatArr[1], floatArr[2], floatArr[3])).toList();
+    }
+
+
     public CharSequence getLine(int lineNumber) {
 
         if (lineNumber >= numberOfLines() || lineNumber < 0) {

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/services/EntityCreationService.java

@@ -247,7 +247,7 @@ public class EntityCreationService {
 
     public void addEntityToGraph(RedactionEntity entity, DocumentTree documentTree) {
 
-        SemanticNode containingNode = documentTree.streamChildrenNodes(Collections.emptyList())
+        SemanticNode containingNode = documentTree.streamChildNodes(Collections.emptyList())
                 .filter(node -> node.buildTextBlock().containsBoundary(entity.getBoundary()))
                 .findFirst()
                 .orElseThrow(() -> new NoSuchElementException("No containing Node found!"));

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/utils/PdfVisualisationUtility.java

@@ -139,7 +139,7 @@ public class PdfVisualisationUtility {
 
     private String buildString(DocumentTree.Entry entry) {
 
-        return entry.getNode().getNumberOnPage() + ": " + entry.getTocId() + ": " + entry.getType().toString();
+        return entry.getNode().getNumberOnPage() + ": " + entry.getTreeId() + ": " + entry.getType().toString();
     }
 
 

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/utils/RectangleMapper.java

@@ -8,14 +8,9 @@ import java.util.List;
 
 import com.iqser.red.service.persistence.service.v1.api.shared.model.redactionlog.Point;
 import com.iqser.red.service.persistence.service.v1.api.shared.model.redactionlog.Rectangle;
-import com.iqser.red.service.redaction.v1.server.layoutparsing.classification.model.AbstractPageBlock;
 
 public class RectangleMapper {
-
+    //TODO: Move functions into using services
-    public static Rectangle2D toRectangle2D(AbstractPageBlock abstractPageBlock) {
-
-        return new Rectangle2D.Float(abstractPageBlock.getMinX(), abstractPageBlock.getMinY(), abstractPageBlock.getWidth(), abstractPageBlock.getHeight());
-    }
 
 
     public static Rectangle2D toRectangle2D(Rectangle redactionLogRectangle) {

redaction-service-v1/redaction-service-server-v1/src/main/java/com/iqser/red/service/redaction/v1/server/layoutparsing/document/utils/RectangleTransformations.java

@@ -29,7 +29,13 @@ public class RectangleTransformations {
 
     public static Rectangle2D bBoxUnionAbstractTextContainer(List<AbstractPageBlock> abstractPageBlocks) {
 
-        return abstractPageBlocks.stream().map(RectangleMapper::toRectangle2D).collect(new Rectangle2DUnion());
+        return abstractPageBlocks.stream().map(RectangleTransformations::toRectangle2D).collect(new Rectangle2DUnion());
+    }
+
+
+    public static Rectangle2D toRectangle2D(AbstractPageBlock abstractPageBlock) {
+
+        return new Rectangle2D.Float(abstractPageBlock.getMinX(), abstractPageBlock.getMinY(), abstractPageBlock.getWidth(), abstractPageBlock.getHeight());
     }