Fix CLI bug

VikParuchuri · Nov 28, 2024 · a580d56 · a580d56
1 parent c8ce327
commit a580d56
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 11 deletions.
diff --git a/extract.py b/extract.py
@@ -1,7 +1,6 @@
 import json
 from collections import defaultdict
 
-import pypdfium2
 import copy
 import os
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tabled-pdf"
-version = "0.1.5"
+version = "0.1.6"
 description = "Detect and recognize tables in PDFs and images."
 authors = ["Vik Paruchuri <[email protected]>"]
 readme = "README.md"

diff --git a/table_app.py b/table_app.py
@@ -104,7 +104,7 @@ def page_count(pdf_file):
     page_number = 1
 
 with col:
-    st.image(pil_image, caption="PDF file (preview)", use_column_width="auto")
+    st.image(pil_image, caption="PDF file (preview)", use_container_width=True)
 
     run_marker = st.sidebar.button("Run Tabled")
 
@@ -122,7 +122,7 @@ def page_count(pdf_file):
 
 for idx, (md, table_img) in enumerate(out_data):
     container.markdown(f"## Table {idx}")
-    container.image(table_img, caption=f"Table {idx}", use_column_width="auto")
+    container.image(table_img, caption=f"Table {idx}", use_container_width=True)
     container.markdown(md)
     container.code(md)
     container.divider()

diff --git a/tabled/inference/detection.py b/tabled/inference/detection.py
@@ -32,10 +32,9 @@ def merge_tables(page_table_boxes):
     return [b for i, b in enumerate(page_table_boxes) if i not in ignore_boxes]
 
 
-def detect_tables(images, highres_images, models, detector_batch_size=settings.DETECTOR_BATCH_SIZE, layout_batch_size=settings.LAYOUT_BATCH_SIZE):
-    det_model, det_processor, layout_model, layout_processor = models
-    line_predictions = batch_text_detection(images, det_model, det_processor, batch_size=detector_batch_size)
-    layout_predictions = batch_layout_detection(images, layout_model, layout_processor, line_predictions, batch_size=layout_batch_size)
+def detect_tables(images, highres_images, models, layout_batch_size=settings.LAYOUT_BATCH_SIZE):
+    layout_model, layout_processor = models
+    layout_predictions = batch_layout_detection(images, layout_model, layout_processor, batch_size=layout_batch_size)
 
     table_imgs = []
     table_counts = []

diff --git a/tabled/inference/models.py b/tabled/inference/models.py
@@ -8,11 +8,9 @@
 
 
 def load_detection_models():
-    det_model = load_det_model()
-    det_processor = load_det_processor()
     layout_model = load_layout_model()
     layout_processor = load_layout_processor()
-    return det_model, det_processor, layout_model, layout_processor
+    return layout_model, layout_processor
 
 
 def load_recognition_models():