feat: small fixes

.github/workflows/build-gpu.yml

@@ -2,6 +2,9 @@ name: Build docker GPU and push ECR
 
 on:
   push:
+    tags:
+      - "v*"
+    branches: [main]
 
 env:
   AWS_REGION: eu-west-1

libs/megaparse/src/megaparse/parser/strategy.py

@@ -11,10 +11,7 @@
 logger = logging.getLogger("megaparse")
 
 
-def get_strategy_page(
-    page: PdfPage,
-    threshold=0.2,
-) -> StrategyEnum:
+def get_strategy_page(page: PdfPage, threshold_image_page: float) -> StrategyEnum:
     total_page_area = page.get_width() * page.get_height()
     total_image_area = 0
     images_coords = []
@@ -43,20 +40,24 @@ def get_strategy_page(
         # Get the total area of the images
     total_image_area = np.sum(canva)
 
-    if total_image_area / total_page_area > threshold:
+    if total_image_area / total_page_area > threshold_image_page:
         return StrategyEnum.HI_RES
     return StrategyEnum.FAST
 
 
-def determine_strategy(file: str | Path | bytes | BinaryIO) -> StrategyEnum:
+def determine_strategy(
+    file: str | Path | bytes | BinaryIO,
+    threshold_pages_ocr: float = 0.2,
+    threshold_image_page: float = 0.4,
+) -> StrategyEnum:
     logger.info("Determining strategy...")
     need_ocr = 0
     document = pdfium.PdfDocument(file)
     for page in document:
-        strategy = get_strategy_page(page)
+        strategy = get_strategy_page(page, threshold_image_page=threshold_image_page)
         need_ocr += strategy == StrategyEnum.HI_RES
 
-    doc_need_ocr = (need_ocr / len(document)) > 0.2
+    doc_need_ocr = (need_ocr / len(document)) > threshold_pages_ocr
     document.close()
 
     if doc_need_ocr:

libs/megaparse/tests/pdf/test_detect_ocr.py

@@ -4,5 +4,7 @@
 
 def test_strategy_all():
     pdf = "/Users/amine/Downloads/RAG Corporate 2024 016.pdf"
-    strategy = determine_strategy(pdf)
+    strategy = determine_strategy(
+        pdf, threshold_pages_ocr=0.2, threshold_image_page=0.3
+    )
     assert strategy == StrategyEnum.HI_RES

libs/megaparse/tests/pdf/test_unstructured_parser.py

@@ -0,0 +1,28 @@
+from pathlib import Path
+
+import pypdfium2 as pdfium
+import pytest
+from megaparse.parser.unstructured_parser import UnstructuredParser
+from megaparse_sdk.schema.parser_config import StrategyEnum
+
+
+@pytest.mark.asyncio
+async def test_unstructured_parser():
+    # scanned pdf
+    p = Path("./tests/pdf/mlbook.pdf")
+    processor = UnstructuredParser(strategy=StrategyEnum.FAST)
+    result = await processor.convert(file_path=p)
+    assert len(result) > 0
+
+
+def test_pdfium():
+    # scanned pdf
+    p = Path("./tests/pdf/mlbook.pdf")
+    document = pdfium.PdfDocument(p)
+
+    objs = []
+    for page in document:
+        for obj in page.get_objects():
+            objs.append(obj)
+
+    document.close()