Browse Source

Merge pull request #2512 from myhloli/dev

fix(ocr): adjust area ratio threshold and update fitz document handling in image conversion
Xiaomeng Zhao 5 months ago
parent
commit
e92b5b698e
2 changed files with 5 additions and 5 deletions
  1. 4 4
      magic_pdf/data/utils.py
  2. 1 1
      magic_pdf/model/batch_analyze.py

+ 4 - 4
magic_pdf/data/utils.py

@@ -10,22 +10,22 @@ from loguru import logger
 
 
 
-def fitz_doc_to_image(doc, dpi=200) -> dict:
+def fitz_doc_to_image(page, dpi=200) -> dict:
     """Convert fitz.Document to image, Then convert the image to numpy array.
 
     Args:
-        doc (_type_): pymudoc page
+        page (_type_): pymudoc page
         dpi (int, optional): reset the dpi of dpi. Defaults to 200.
 
     Returns:
         dict:  {'img': numpy array, 'width': width, 'height': height }
     """
     mat = fitz.Matrix(dpi / 72, dpi / 72)
-    pm = doc.get_pixmap(matrix=mat, alpha=False)
+    pm = page.get_pixmap(matrix=mat, alpha=False)
 
     # If the width or height exceeds 4500 after scaling, do not scale further.
     if pm.width > 4500 or pm.height > 4500:
-        pm = doc.get_pixmap(matrix=fitz.Matrix(1, 1), alpha=False)
+        pm = page.get_pixmap(matrix=fitz.Matrix(1, 1), alpha=False)
 
     # Convert pixmap samples directly to numpy array
     img = np.frombuffer(pm.samples, dtype=np.uint8).reshape(pm.height, pm.width, 3)

+ 1 - 1
magic_pdf/model/batch_analyze.py

@@ -156,7 +156,7 @@ class BatchAnalyze:
                         res_area = get_coords_and_area(res)[4]
                         if res_area > 0:
                             ratio = ocr_res_area / res_area
-                            if ratio > 0.45:
+                            if ratio > 0.25:
                                 res["category_id"] = 1
                             else:
                                 continue