Upload folder using huggingface_hub

app.py

@@ -6,9 +6,10 @@
 import cv2
 import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 import streamlit as st
 import torch
-from backend.pytorch import DET_ARCHS, RECO_ARCHS, forward_image, load_predictor
+from backend.pytorch import DET_ARCHS, LAYOUT_ARCHS, RECO_ARCHS, forward_image, load_predictor
 
 from doctr.io import DocumentFile
 from doctr.utils.visualization import visualize_page
@@ -16,7 +17,7 @@ from doctr.utils.visualization import visualize_page
 forward_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 
-def main(det_archs, reco_archs):
+def main(det_archs, reco_archs, layout_archs):
     """Build a streamlit layout"""
     # Wide mode
     st.set_page_config(layout="wide")
@@ -67,6 +68,11 @@ def main(det_archs, reco_archs):
     straighten_pages = st.sidebar.checkbox("Straighten pages", value=False)
     # Export as straight boxes
     export_straight_boxes = st.sidebar.checkbox("Export as straight boxes", value=False)
+    # Layout detection
+    detect_layout = st.sidebar.checkbox("Detect layout", value=False)
+    layout_arch = st.sidebar.selectbox("Layout detection model", layout_archs, disabled=not detect_layout)
+    # Table detection (relies on the layout model to locate tables)
+    detect_tables = st.sidebar.checkbox("Detect tables", value=False)
     st.sidebar.write("\n")
     # Binarization threshold
     bin_thresh = st.sidebar.slider("Binarization threshold", min_value=0.1, max_value=0.9, value=0.3, step=0.1)
@@ -92,6 +98,9 @@ def main(det_archs, reco_archs):
                     bin_thresh=bin_thresh,
                     box_thresh=box_thresh,
                     device=forward_device,
+                    detect_layout=detect_layout,
+                    layout_arch=layout_arch,
+                    detect_tables=detect_tables,
                 )
 
             with st.spinner("Analyzing..."):
@@ -117,10 +126,17 @@ def main(det_archs, reco_archs):
                     img = out.pages[0].synthesize()
                     cols[3].image(img, clamp=True)
 
+                # Display extracted tables (if any)
+                if out.pages[0].tables:
+                    st.markdown("\nExtracted tables:")
+                    for idx, table in enumerate(out.pages[0].tables):
+                        st.markdown(f"**Table {idx + 1}** ({table.num_rows} x {table.num_cols})")
+                        st.dataframe(pd.DataFrame(table.to_grid()))
+
                 # Display JSON
                 st.markdown("\nHere are your analysis results in JSON format:")
                 st.json(page_export, expanded=False)
 
 
 if __name__ == "__main__":
-    main(DET_ARCHS, RECO_ARCHS)
+    main(DET_ARCHS, RECO_ARCHS, LAYOUT_ARCHS)

backend/pytorch.py

@@ -31,6 +31,10 @@ RECO_ARCHS = [
     "parseq",
     "viptr_tiny",
 ]
+LAYOUT_ARCHS = [
+    "lw_detr_s",
+    "lw_detr_m",
+]
 
 
 def load_predictor(
@@ -44,6 +48,9 @@ def load_predictor(
     bin_thresh: float,
     box_thresh: float,
     device: torch.device,
+    detect_layout: bool,
+    layout_arch: str,
+    detect_tables: bool,
 ) -> OCRPredictor:
     """Load a predictor from doctr.models
 
@@ -58,6 +65,9 @@ def load_predictor(
         bin_thresh: binarization threshold for the segmentation map
         box_thresh: minimal objectness score to consider a box
         device: torch.device, the device to load the predictor on
+        detect_layout: whether to run a layout detection model and attach the regions to each page
+        layout_arch: layout architecture to use when detect_layout is True
+        detect_tables: whether to detect tables (via the layout model), structure them and attach them to each page
 
     Returns:
         instance of OCRPredictor
@@ -72,6 +82,9 @@ def load_predictor(
         detect_orientation=not assume_straight_pages,
         disable_page_orientation=disable_page_orientation,
         disable_crop_orientation=disable_crop_orientation,
+        detect_layout=detect_layout,
+        layout_arch=layout_arch,
+        detect_tables=detect_tables,
     ).to(device)
     predictor.det_predictor.model.postprocessor.bin_thresh = bin_thresh
     predictor.det_predictor.model.postprocessor.box_thresh = box_thresh

requirements.txt

@@ -1,2 +1,3 @@
 -e "python-doctr[viz] @ git+https://github.com/mindee/doctr.git"
 streamlit>=1.0.0
+pandas>=2.0.0