Create detect_layout.py

detect_layout.py

@@ -0,0 +1,34 @@
+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, load_processor
+from surya.postprocessing.heatmap import draw_polys_on_image
+from surya.settings import settings
+import os
+
+def main(input_path, max_pages=None):
+
+    model = load_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
+    processor = load_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
+    det_model = load_model()
+    det_processor = load_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, model, processor, line_predictions)
+
+
+    for idx, (image, layout_pred, name) in enumerate(zip(images, layout_predictions, names)):
+        polygons = [p.polygon for p in layout_pred.bboxes]
+        labels = [p.label for p in layout_pred.bboxes]
+        bbox_image = draw_polys_on_image(polygons, copy.deepcopy(image), labels=labels)
+        return bbox_image