diff --git a/pix2text/table_ocr.py b/pix2text/table_ocr.py index 31e3cfa..805474e 100644 --- a/pix2text/table_ocr.py +++ b/pix2text/table_ocr.py @@ -157,6 +157,22 @@ def recognize( out_markdown=True, **kwargs, ) -> Dict[str, Any]: + """ + + Args: + img (): + tokens (): + out_objects (): + out_cells (): + out_html (): + out_csv (): + out_markdown (): + **kwargs (): + * save_analysis_res (str): Save the parsed result image in this file; default value is `None`, which means not to save + + Returns: + + """ out_formats = {} if self.str_model is None: print("No structure model loaded.") @@ -196,6 +212,9 @@ def recognize( self._ocr_texts(img, cells) if out_cells: out_formats['cells'] = tables_cells + if kwargs.get('save_analysis_res'): + visualize_cells(img, tables_cells[0], kwargs['save_analysis_res']) + if not (out_html or out_csv): return out_formats diff --git a/tests/test_table_ocr.py b/tests/test_table_ocr.py index 66b6d24..6ba7e5e 100644 --- a/tests/test_table_ocr.py +++ b/tests/test_table_ocr.py @@ -2,9 +2,8 @@ import pytest import os -from pix2text.utils import read_img from pix2text.ocr_engine import prepare_ocr_engine -from pix2text.table_ocr import TableOCR, visualize_cells +from pix2text.table_ocr import TableOCR def test_recognize(): @@ -13,9 +12,16 @@ def test_recognize(): languages = ('en', 'ch_sim') text_ocr = prepare_ocr_engine(languages, {}) ocr = TableOCR(text_ocr=text_ocr) - result = ocr.recognize(image_path, out_csv=True, out_cells=True, out_objects=False, out_html=True, out_markdown=True) + result = ocr.recognize( + image_path, + out_csv=True, + out_cells=True, + out_objects=False, + out_html=True, + out_markdown=True, + save_analysis_res='out-table-rec.png', + ) - visualize_cells(read_img(image_path, 'Image'), result['cells'][0], 'out-table-rec.png') print(result) @@ -25,7 +31,14 @@ def test_recognize2(): languages = ('en', 'ch_sim') text_ocr = prepare_ocr_engine(languages, {}) ocr = TableOCR.from_config(text_ocr=text_ocr) - result = ocr.recognize(image_path, out_csv=True, out_cells=True, out_objects=False, out_html=True, out_markdown=True) + result = ocr.recognize( + image_path, + out_csv=True, + out_cells=True, + out_objects=False, + out_html=True, + out_markdown=True, + save_analysis_res='out-table-rec.png', + ) - visualize_cells(read_img(image_path, 'Image'), result['cells'][0], 'out-table-rec.png') print(result)