Create reading_order.py

reading_order.py

@@ -0,0 +1,45 @@
+import os
+import copy
+
+from surya.detection import batch_text_detection
+from surya.input.load import load_from_folder, load_from_file
+from surya.layout import batch_layout_detection
+from surya.model.detection.model import load_model as load_det_model, load_processor as load_det_processor
+from surya.model.ordering.model import load_model
+from surya.model.ordering.processor import load_processor
+from surya.ordering import batch_ordering
+from surya.postprocessing.heatmap import draw_polys_on_image
+from surya.settings import settings
+
+
+def main(input_path, max_pages=None):
+    model = load_model()
+    processor = load_processor()
+
+    layout_model = load_det_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
+    layout_processor = load_det_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
+
+    det_model = load_det_model()
+    det_processor = load_det_processor()
+
+    if os.path.isdir(input_path):
+        images, names = load_from_folder(input_path, max_pages)
+
+    else:
+        images, names = load_from_file(input_path, max_pages)
+
+
+    line_predictions = batch_text_detection(images, det_model, det_processor)
+    layout_predictions = batch_layout_detection(images, layout_model, layout_processor, line_predictions)
+    bboxes = []
+    for layout_pred in layout_predictions:
+        bbox = [l.bbox for l in layout_pred.bboxes]
+        bboxes.append(bbox)
+
+    order_predictions = batch_ordering(images, bboxes, model, processor)
+
+    for idx, (image, layout_pred, order_pred, name) in enumerate(zip(images, layout_predictions, order_predictions, names)):
+        polys = [l.polygon for l in order_pred.bboxes]
+        labels = [str(l.position) for l in order_pred.bboxes]
+        bbox_image = draw_polys_on_image(polys, copy.deepcopy(image), labels=labels, label_font_size=20)
+        return bbox_image