New layout model

VikParuchuri · Dec 30, 2024 · ffd3f5f · ffd3f5f
1 parent f8188f4
commit ffd3f5f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/ocr_app.py b/ocr_app.py
@@ -84,7 +84,7 @@ def table_recognition(img, highres_img, skip_table_detection: bool) -> (Image.Im
         labels = []
         colors = []
 
-        for item in results.rows + results.cols:
+        for item in results.cells:
             adjusted_bboxes.append([
                 (item.bbox[0] + table_bbox[0]),
                 (item.bbox[1] + table_bbox[1]),

diff --git a/surya/settings.py b/surya/settings.py
@@ -65,7 +65,7 @@ def TORCH_DEVICE_MODEL(self) -> str:
     RECOGNITION_ENCODER_BATCH_DIVISOR: int = 1 # Divisor for batch size in decoder
 
     # Layout
-    LAYOUT_MODEL_CHECKPOINT: str = "datalab-to/surya_layout0"
+    LAYOUT_MODEL_CHECKPOINT: str = "datalab-to/surya_layout"
     LAYOUT_IMAGE_SIZE: Dict = {"height": 768, "width": 768}
     LAYOUT_SLICE_MIN: Dict = {"height": 1500, "width": 1500} # When to start slicing images
     LAYOUT_SLICE_SIZE: Dict = {"height": 1200, "width": 1200} # Size of slices