Resort bboxes based on layout

Author: VikParuchuri

ocr_app.py

@@ -70,7 +70,8 @@ def layout_detection(img) -> (Image.Image, LayoutResult):
 def order_detection(img) -> (Image.Image, OrderResult):
     _, layout_pred = layout_detection(img)
     bboxes = [l.bbox for l in layout_pred.bboxes]
-    pred = batch_ordering([img], [bboxes], order_model, order_processor)[0]
+    labels = [l.label for l in layout_pred.bboxes]
+    pred = batch_ordering([img], [bboxes], order_model, order_processor, labels=[labels])[0]
     polys = [l.polygon for l in pred.bboxes]
     positions = [str(l.position) for l in pred.bboxes]
     order_img = draw_polys_on_image(polys, img.copy(), labels=positions, label_font_size=20)

reading_order.py

@@ -42,11 +42,14 @@ def main():
     line_predictions = batch_text_detection(images, det_model, det_processor)
     layout_predictions = batch_layout_detection(images, layout_model, layout_processor, line_predictions)
     bboxes = []
+    labels = []
     for layout_pred in layout_predictions:
         bbox = [l.bbox for l in layout_pred.bboxes]
+        label = [l.label for l in layout_pred.bboxes]
         bboxes.append(bbox)
+        labels.append(label)
 
-    order_predictions = batch_ordering(images, bboxes, model, processor)
+    order_predictions = batch_ordering(images, bboxes, model, processor, labels=labels)
     result_path = os.path.join(args.results_dir, folder_name)
     os.makedirs(result_path, exist_ok=True)
 

surya/ordering.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import List
+from typing import List, Optional
 import torch
 from PIL import Image
 
@@ -30,11 +30,16 @@ def rank_elements(arr):
     return rank
 
 
-def batch_ordering(images: List, bboxes: List[List[List[float]]], model, processor) -> List[OrderResult]:
+def batch_ordering(images: List, bboxes: List[List[List[float]]], model, processor, labels: Optional[List[List[str]]] = None) -> List[OrderResult]:
     assert all([isinstance(image, Image.Image) for image in images])
     assert len(images) == len(bboxes)
     batch_size = get_batch_size()
 
+    if labels is not None:
+        assert len(labels) == len(images)
+        for l, b in zip(labels, bboxes):
+            assert len(l) == len(b)
+
     images = [image.convert("RGB") for image in images]
 
     output_order = []
@@ -64,9 +69,29 @@ def batch_ordering(images: List, bboxes: List[List[List[float]]], model, process
         for j in range(logits.shape[0]):
             row_logits = logits[j].tolist()
             row_bboxes = bboxes[i+j]
+            assert len(row_logits) == len(row_bboxes), "Mismatch between logits and bboxes."
+
             orig_size = orig_sizes[j]
             ranks = rank_elements(row_logits)
 
+            if labels is not None:
+                # This is to force headers/footers into the proper order
+                row_label = labels[i+j]
+                combined = [[i, bbox, label, rank] for i, (bbox, label, rank) in enumerate(zip(row_bboxes, row_label, ranks))]
+                combined = sorted(combined, key=lambda x: x[3])
+
+                sorted_boxes = ([row for row in combined if row[2] == "Page-header"] +
+                                [row for row in combined if row[2] not in ["Page-header", "Page-footer"]] +
+                                [row for row in combined if row[2] == "Page-footer"])
+
+                # Re-rank after sorting
+                for rank, row in enumerate(sorted_boxes):
+                    row[3] = rank
+
+                sorted_boxes = sorted(sorted_boxes, key=lambda x: x[0])
+                row_bboxes = [row[1] for row in sorted_boxes]
+                ranks = [row[3] for row in sorted_boxes]
+
             order_boxes = []
             for row_bbox, rank in zip(row_bboxes, ranks):
                 order_box = OrderBox(

surya/settings.py

@@ -68,11 +68,11 @@ def TORCH_DEVICE_DETECTION(self) -> str:
     RECOGNITION_PAD_VALUE: int = 0 # Should be 0 or 255
 
     # Layout
-    LAYOUT_MODEL_CHECKPOINT: str = "vikp/surya_layout"
+    LAYOUT_MODEL_CHECKPOINT: str = "vikp/surya_layout2"
     LAYOUT_BENCH_DATASET_NAME: str = "vikp/publaynet_bench"
 
     # Ordering
-    ORDER_MODEL_CHECKPOINT: str = "vikp/order2"
+    ORDER_MODEL_CHECKPOINT: str = "vikp/surya_order"
     ORDER_IMAGE_SIZE: Dict = {"height": 1024, "width": 1024}
     ORDER_MAX_BOXES: int = 256
     ORDER_BATCH_SIZE: Optional[int] = None  # Defaults to 4 for CPU/MPS, 32 otherwise