Update rapid_table to v0.1.2

RapidAI · Dec 27, 2023 · 35ae438 · 35ae438
1 parent b800b15
commit 35ae438
Show file tree

Hide file tree

Showing 11 changed files with 143 additions and 150 deletions.
diff --git a/demo.py b/demo.py
@@ -9,7 +9,7 @@
 
 from rapid_layout import RapidLayout
 from rapid_orientation import RapidOrientation
-from rapid_table import RapidTable
+from rapid_table import RapidTable, VisTable
 
 
 def vis_layout(img: np.ndarray, layout_res: list) -> None:
@@ -35,24 +35,6 @@ def vis_layout(img: np.ndarray, layout_res: list) -> None:
     print(f"The infer result has saved in {image_save}")
 
 
-def vis_table(table_res):
-    style_res = """<style>td {border-left: 1px solid;border-bottom:1px solid;}
-                   table, th {border-top:1px solid;font-size: 10px;
-                   border-collapse: collapse;border-right: 1px solid;}
-                </style>"""
-    prefix_table, suffix_table = table_res.split("<body>")
-    new_table_res = f"{prefix_table}{style_res}<body>{suffix_table}"
-
-    draw_img_save = Path("./inference_results/")
-    if not draw_img_save.exists():
-        draw_img_save.mkdir(parents=True, exist_ok=True)
-
-    html_path = str(draw_img_save / "table_result.html")
-    with open(html_path, "w", encoding="utf-8") as f:
-        f.write(new_table_res)
-    print(f"The infer result has saved in {html_path}")
-
-
 def demo_layout():
     layout_engine = RapidLayout()
 
@@ -69,12 +51,17 @@ def demo_table():
 
     table_engine = RapidTable()
     ocr_engine = RapidOCR()
+    viser = VisTable()
 
     img_path = "tests/test_files/table.jpg"
     ocr_result, _ = ocr_engine(img_path)
-    table_html_str, _ = table_engine(img_path, ocr_result)
+    table_html_str, table_cell_bboxes, _ = table_engine(img_path, ocr_result)
+
+    save_dir = Path("./inference_results/")
+    if not save_dir.exists():
+        save_dir.mkdir(parents=True, exist_ok=True)
 
-    vis_table(table_html_str)
+    viser(img_path, save_dir, table_html_str, table_cell_bboxes)
     print(table_html_str)
 
 
@@ -86,6 +73,6 @@ def demo_orientation():
 
 
 if __name__ == "__main__":
-    demo_layout()
-    # demo_table()
+    # demo_layout()
+    demo_table()
     # demo_orientation()
diff --git a/docs/README_Table.md b/docs/README_Table.md
@@ -11,22 +11,27 @@
 - 具体来说，就是分析给定的表格图像，将表格还原为对应的html格式。
 - 目前支持两种类别的表格识别模型：中文和英文表格识别模型，具体可参见下面表格：
 
-    | 模型类型  |        模型名称         | 模型大小 |
-    |:---:|:---:|:---:|
-    |   英文   | `en_ppstructure_mobile_v2_SLANet.onnx` |  7.3M   |
-    |   中文   |   `ch_ppstructure_mobile_v2_SLANet.onnx`    |  7.4M   |
-- 模型下载地址为：[百度网盘](https://pan.baidu.com/s/1PI9fksW6F6kQfJhwUkewWg?pwd=p29g) | [Google Drive](https://drive.google.com/drive/folders/1DAPWSN2zGQ-ED_Pz7RaJGTjfkN2-Mvsf?usp=sharing)
+    | 模型类型 |                模型名称                | 模型大小 |
+    | :------: | :------------------------------------: | :------: |
+    |   英文   | `en_ppstructure_mobile_v2_SLANet.onnx` |   7.3M   |
+    |   中文   | `ch_ppstructure_mobile_v2_SLANet.onnx` |   7.4M   |
 
 
+- 模型下载地址为：[百度网盘](https://pan.baidu.com/s/1PI9fksW6F6kQfJhwUkewWg?pwd=p29g) | [Google Drive](https://drive.google.com/drive/folders/1DAPWSN2zGQ-ED_Pz7RaJGTjfkN2-Mvsf?usp=sharing) |
+
+
+#### 安装
+由于模型较小，预先将英文表格识别模型(`en_ppstructure_mobile_v2_SLANet.onnx`)打包进了whl包内，如果做英文表格识别，可直接安装使用
+
+> ⚠️注意：`rapid_table>=v0.1.0`之后，不再将`rapidocr_onnxruntime`依赖强制打包到`rapid_table`中。使用前，需要自行安装`rapidocr_onnxruntime`包。
+
+```bash
+pip install rapidocr_onnxruntime
+pip install rapid_table
+```
+
 #### 使用方式
-1. pip安装
-    - 由于模型较小，预先将英文表格识别模型(`en_ppstructure_mobile_v2_SLANet.onnx`)打包进了whl包内，如果做英文表格识别，可直接安装使用
-    - ⚠️注意：`rapid_table>=v0.1.0`之后，不再将`rapidocr_onnxruntime`依赖强制打包到`rapid_table`中。使用前，需要自行安装`rapidocr_onnxruntime`包。
-        ```bash
-        pip install rapidocr_onnxruntime
-        pip install rapid_table
-        ```
-2. python脚本运行
+1. python脚本运行
     ````python
     from rapid_table import RapidTable
     from rapidocr_onnxruntime import RapidOCR
@@ -39,19 +44,25 @@
     img_path = 'test_images/table.jpg'
 
     ocr_result, _ = ocr_engine(img_path)
-    table_html_str, _ = table_engine(img_path, ocr_result)
+    table_html_str, table_cell_bboxes, elapse = table_engine(img_path, ocr_result)
 
+    save_dir = Path("./inference_results/")
+    if not save_dir.exists():
+        save_dir.mkdir(parents=True, exist_ok=True)
+
+    # 可视化结果，如果不提供table_cell_bboxes，只会导出html文件
+    viser(img_path, save_dir, table_html_str, table_cell_bboxes)
     print(table_html_str)
     ````
-3. 终端运行
+2. 终端运行
    - 用法:
      ```bash
      $ rapid_table -h
      usage: rapid_table [-h] [-v] -img IMG_PATH [-m MODEL_PATH]
 
      optional arguments:
      -h, --help            show this help message and exit
-     -v, --vis             Wheter to visualize the layout results.
+     -v, --vis             Whether to visualize the layout results.
      -img IMG_PATH, --img_path IMG_PATH
                            Path to image for layout.
      -m MODEL_PATH, --model_path MODEL_PATH
@@ -62,7 +73,7 @@
      $ rapid_table -v -img test_images/table.jpg
      ```
 
-4. 结果
+3. 结果
     - 返回结果
         ```html
         <html><body><table><tr><td>Methods</td><td></td><td></td><td></td><td>FPS</td></tr><tr><td>SegLink [26]</td><td>70.0</td><td>86d><td.0</td><td>77.0</td><td>8.9</td></tr><tr><td>PixelLink [4]</td><td>73.2</td><td>83.0</td><td>77.8</td><td></td></tr><tr><td>TextSnake [18]</td><td>73.9</td><td>83.2</td><td>78.3</td><td>1.1</td></tr><tr><td>TextField [37]</td><td>75.9</td><td>87.4</td><td>81.3</td><td>5.2</td></tr><tr><td>MSR[38]</td><td>76.7</td><td>87.87.4</td><td>81.7</td><td></td></tr><tr><td>FTSN [3]</td><td>77.1</td><td>87.6</td><td>82.0</td><td></td></tr><tr><td>LSE[30]</td><td>81.7</td><td>84.2</td><td>82.9</td><><ttd></td></tr><tr><td>CRAFT [2]</td><td>78.2</td><td>88.2</td><td>82.9</td><td>8.6</td></tr><tr><td>MCN[16]</td><td>79</td><td>88</td><td>83</td><td></td></tr><tr><td>ATRR</>[35]</td><td>82.1</td><td>85.2</td><td>83.6</td><td></td></tr><tr><td>PAN [34]</td><td>83.8</td><td>84.4</td><td>84.1</td><td>30.2</td></tr><tr><td>DB[12]</td><td>79.2</t91/d><td>91.5</td><td>84.9</td><td>32.0</td></tr><tr><td>DRRG[41]</td><td>82.30</td><td>88.05</td><td>85.08</td><td></td></tr><tr><td>Ours (SynText)</td><td>80.68</td><td>85<t..40</td><td>82.97</td><td>12.68</td></tr><tr><td>Ours (MLT-17)</td><td>84.54</td><td>86.62</td><td>85.57</td><td>12.31</td></tr></table></body></html>
@@ -73,14 +84,21 @@
     </div>
 
 #### 更新日志
-- 2023-07-17 v0.1.0 update:
-    - 将`rapidocr_onnxruntime`部分从`rapid_table`中解耦合出来，给出选项是否依赖，更加灵活。
-    - 增加接口输入参数`ocr_result`：
-      - 如果在调用函数时，事先指定了`ocr_result`参数值，则不会再走OCR。其中`ocr_result`格式需要和`rapidocr_onnxruntime`返回值一致。
-      - 如果未指定`ocr_result`参数值，但是事先安装了`rapidocr_onnxruntime`库，则会自动调用该库，进行识别。
-      - 如果`ocr_result`未指定，且`rapidocr_onnxruntime`未安装，则会报错。必须满足两个条件中一个。
-- 2023-07-10 v0.0.13 updata:
-    - 更改传入表格还原中OCR的实例接口，可以传入其他OCR实例，前提要与`rapidocr_onnxruntime`接口一致
-- 2023-07-06 v0.0.12 update:
-    - 去掉返回表格的html字符串中的`<thead></thead><tbody></tbody>`元素，便于后续统一。
-    - 采用Black工具优化代码
+##### 2023-12-27 v0.1.2 update:
+- 添加返回cell坐标框参数
+- 完善可视化函数
+
+##### 2023-07-17 v0.1.0 update:
+- 将`rapidocr_onnxruntime`部分从`rapid_table`中解耦合出来，给出选项是否依赖，更加灵活。
+
+- 增加接口输入参数`ocr_result`：
+  - 如果在调用函数时，事先指定了`ocr_result`参数值，则不会再走OCR。其中`ocr_result`格式需要和`rapidocr_onnxruntime`返回值一致。
+  - 如果未指定`ocr_result`参数值，但是事先安装了`rapidocr_onnxruntime`库，则会自动调用该库，进行识别。
+  - 如果`ocr_result`未指定，且`rapidocr_onnxruntime`未安装，则会报错。必须满足两个条件中一个。
+
+##### 2023-07-10 v0.0.13 updata:
+- 更改传入表格还原中OCR的实例接口，可以传入其他OCR实例，前提要与`rapidocr_onnxruntime`接口一致
+
+##### 2023-07-06 v0.0.12 update:
+- 去掉返回表格的html字符串中的`<thead></thead><tbody></tbody>`元素，便于后续统一。
+- 采用Black工具优化代码
diff --git a/docs/doc_whl_rapid_table.md b/docs/doc_whl_rapid_table.md
@@ -1,67 +1 @@
-## rapid-table
-<p align="left">
-    <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.12-aff.svg"></a>
-    <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
-    <a href="https://pypi.org/project/rapid-table/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid-table"></a>
-    <a href="https://pepy.tech/project/rapid-table"><img src="https://static.pepy.tech/personalized-badge/rapid-table?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></a>
-</p>
-
-
-### 1. Install package by pypi.
-⚠️Attention: After `rapid_table>=v0.1.0`, you need to install `rapidocr_onnxruntime` package firstly.
-```bash
-pip install rapidocr_onnxruntime
-pip install rapid-table
-```
-
-### 2. Run by script.
-- RapidTable has the default `model_path` value, you can set the different value of `model_path` to use different models, e.g. `table_engine = RapidTable(model_path='ch_ppstructure_mobile_v2_SLANet.onnx')`
-- See details, for [README_Table](https://github.com/RapidAI/RapidStructure/blob/main/docs/README_Table.md) .
-- 📌 `table.jpg` source: [link](https://github.com/RapidAI/RapidStructure/blob/main/test_images/table.jpg)
-
-    ````python
-    from rapid_table import RapidTable
-    from rapidocr_onnxruntime import RapidOCR
-
-    table_engine = RapidTable()
-    ocr_engine = RapidOCR()
-
-    img_path = 'test_images/table.jpg'
-
-    ocr_result, _ = ocr_engine(img_path)
-    table_html_str, _ = table_engine(img_path, ocr_result)
-
-    print(table_html_str)
-    ````
-
-### 3. Run by command line.
-- Usage:
-    ```bash
-    $ rapid_table -h
-    usage: rapid_table [-h] [-v] -img IMG_PATH [-m MODEL_PATH]
-
-    optional arguments:
-    -h, --help            show this help message and exit
-    -v, --vis             Wheter to visualize the layout results.
-    -img IMG_PATH, --img_path IMG_PATH
-                        Path to image for layout.
-    -m MODEL_PATH, --model_path MODEL_PATH
-                        The model path used for inference.
-    ```
-
-- Example:
-    ```bash
-    $ rapid_table -v -img test_images/table.jpg
-    ```
-
-### 4. Result.
-- Return value.
-    ```html
-    <html><body><table><tr><td>Methods</td><td></td><td></td><td></td><td>FPS</td></tr><tr><td>SegLink [26]</td><td>70.0</td><td>86d><td.0</td><td>77.0</td><td>8.9</td></tr><tr><td>PixelLink [4]</td><td>73.2</td><td>83.0</td><td>77.8</td><td></td></tr><tr><td>TextSnake [18]</td><td>73.9</td><td>83.2</td><td>78.3</td><td>1.1</td></tr><tr><td>TextField [37]</td><td>75.9</td><td>87.4</td><td>81.3</td><td>5.2</td></tr><tr><td>MSR[38]</td><td>76.7</td><td>87.87.4</td><td>81.7</td><td></td></tr><tr><td>FTSN [3]</td><td>77.1</td><td>87.6</td><td>82.0</td><td></td></tr><tr><td>LSE[30]</td><td>81.7</td><td>84.2</td><td>82.9</td><><ttd></td></tr><tr><td>CRAFT [2]</td><td>78.2</td><td>88.2</td><td>82.9</td><td>8.6</td></tr><tr><td>MCN[16]</td><td>79</td><td>88</td><td>83</td><td></td></tr><tr><td>ATRR</>[35]</td><td>82.1</td><td>85.2</td><td>83.6</td><td></td></tr><tr><td>PAN [34]</td><td>83.8</td><td>84.4</td><td>84.1</td><td>30.2</td></tr><tr><td>DB[12]</td><td>79.2</t91/d><td>91.5</td><td>84.9</td><td>32.0</td></tr><tr><td>DRRG[41]</td><td>82.30</td><td>88.05</td><td>85.08</td><td></td></tr><tr><td>Ours (SynText)</td><td>80.68</td><td>85<t..40</td><td>82.97</td><td>12.68</td></tr><tr><td>Ours (MLT-17)</td><td>84.54</td><td>86.62</td><td>85.57</td><td>12.31</td></tr></table></body></html>
-    ```
-- Visualize result.
-    <div align="center">
-        <table><tr><td>Methods</td><td></td><td></td><td></td><td>FPS</td></tr><tr><td>SegLink [26]</td><td>70.0</td><td>86d><td.0</td><td>77.0</td><td>8.9</td></tr><tr><td>PixelLink [4]</td><td>73.2</td><td>83.0</td><td>77.8</td><td></td></tr><tr><td>TextSnake [18]</td><td>73.9</td><td>83.2</td><td>78.3</td><td>1.1</td></tr><tr><td>TextField [37]</td><td>75.9</td><td>87.4</td><td>81.3</td><td>5.2</td></tr><tr><td>MSR[38]</td><td>76.7</td><td>87.87.4</td><td>81.7</td><td></td></tr><tr><td>FTSN [3]</td><td>77.1</td><td>87.6</td><td>82.0</td><td></td></tr><tr><td>LSE[30]</td><td>81.7</td><td>84.2</td><td>82.9</td><><ttd></td></tr><tr><td>CRAFT [2]</td><td>78.2</td><td>88.2</td><td>82.9</td><td>8.6</td></tr><tr><td>MCN[16]</td><td>79</td><td>88</td><td>83</td><td></td></tr><tr><td>ATRR</>[35]</td><td>82.1</td><td>85.2</td><td>83.6</td><td></td></tr><tr><td>PAN [34]</td><td>83.8</td><td>84.4</td><td>84.1</td><td>30.2</td></tr><tr><td>DB[12]</td><td>79.2</t91/d><td>91.5</td><td>84.9</td><td>32.0</td></tr><tr><td>DRRG[41]</td><td>82.30</td><td>88.05</td><td>85.08</td><td></td></tr><tr><td>Ours (SynText)</td><td>80.68</td><td>85<t..40</td><td>82.97</td><td>12.68</td></tr><tr><td>Ours (MLT-17)</td><td>84.54</td><td>86.62</td><td>85.57</td><td>12.31</td></tr></table>
-    </div>
-
 ### For details, see [Rapid Table](https://github.com/RapidAI/RapidStructure/blob/main/docs/README_Table.md)
diff --git a/rapid_table/__init__.py b/rapid_table/__init__.py
@@ -1,4 +1,5 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: [email protected]
-from .rapid_table import RapidTable
+from .main import RapidTable
+from .utils import VisTable
diff --git a/rapid_table/rapid_table.py → rapid_table/main.py b/rapid_table/rapid_table.py → rapid_table/main.py
@@ -13,7 +13,7 @@
 
 from .table_matcher import TableMatch
 from .table_structure import TableStructurer
-from .utils import LoadImage, vis_table
+from .utils import LoadImage, VisTable
 
 root_dir = Path(__file__).resolve().parent
 
@@ -56,11 +56,11 @@ def __call__(
             ocr_result, _ = self.ocr_engine(img)
         dt_boxes, rec_res = self.get_boxes_recs(ocr_result, h, w)
 
-        structure_res, _ = self.table_structure(copy.deepcopy(img))
-        pred_html = self.table_matcher(structure_res, dt_boxes, rec_res)
+        pred_structures, pred_bboxes, _ = self.table_structure(copy.deepcopy(img))
+        pred_html = self.table_matcher(pred_structures, pred_bboxes, dt_boxes, rec_res)
 
         elapse = time.time() - s
-        return pred_html, elapse
+        return pred_html, pred_bboxes, elapse
 
     def get_boxes_recs(
         self, ocr_result: List[Union[List[List[float]], str, str]], h: int, w: int
@@ -113,13 +113,14 @@ def main() -> None:
     img = cv2.imread(args.img_path)
 
     ocr_result, _ = ocr_engine(img)
-    table_html_str, elapse = rapid_table(img, ocr_result)
+    table_html_str, table_cell_bboxes, elapse = rapid_table(img, ocr_result)
     print(table_html_str)
 
+    viser = VisTable()
     if args.vis:
         img_path = Path(args.img_path)
-        save_path = img_path.resolve().parent / f"vis_{img_path.stem}.html"
-        vis_table(table_html_str, str(save_path))
+        save_dir = img_path.resolve().parent
+        viser(img_path, save_dir, table_html_str, table_cell_bboxes)
 
 
 if __name__ == "__main__":

diff --git a/rapid_table/table_matcher/matcher.py b/rapid_table/table_matcher/matcher.py
@@ -22,8 +22,7 @@ def __init__(self, filter_ocr_result=True, use_master=False):
         self.filter_ocr_result = filter_ocr_result
         self.use_master = use_master
 
-    def __call__(self, structure_res, dt_boxes, rec_res):
-        pred_structures, pred_bboxes = structure_res
+    def __call__(self, pred_structures, pred_bboxes, dt_boxes, rec_res):
         if self.filter_ocr_result:
             dt_boxes, rec_res = self._filter_ocr_result(pred_bboxes, dt_boxes, rec_res)
         matched_index = self.match_result(dt_boxes, pred_bboxes)

diff --git a/rapid_table/table_structure/table_structure.py b/rapid_table/table_structure/table_structure.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import time
+
 import numpy as np
 
 from .utils import OrtInferSession, TableLabelDecode, TablePreprocess
@@ -53,4 +54,4 @@ def __call__(self, img):
             + ["</table>", "</body>", "</html>"]
         )
         elapse = time.time() - starttime
-        return (structure_str_list, bbox_list), elapse
+        return structure_str_list, bbox_list, elapse
diff --git a/rapid_table/table_structure/utils.py b/rapid_table/table_structure/utils.py
@@ -158,14 +158,17 @@ def decode_label(self, batch):
                 char_idx = int(structure_idx[batch_idx][idx])
                 if idx > 0 and char_idx == end_idx:
                     break
+
                 if char_idx in ignored_tokens:
                     continue
+
                 structure_list.append(self.character[char_idx])
 
                 bbox = gt_bbox_list[batch_idx][idx]
                 if bbox.sum() != 0:
                     bbox = self._bbox_decode(bbox, shape_list[batch_idx])
                     bbox_list.append(bbox)
+
             structure_batch_list.append(structure_list)
             bbox_batch_list.append(bbox_list)
         result = {
@@ -187,12 +190,12 @@ def get_ignored_tokens(self):
 
     def get_beg_end_flag_idx(self, beg_or_end):
         if beg_or_end == "beg":
-            idx = np.array(self.dict[self.beg_str])
-        elif beg_or_end == "end":
-            idx = np.array(self.dict[self.end_str])
-        else:
-            assert False, "unsupport type %s in get_beg_end_flag_idx" % beg_or_end
-        return idx
+            return np.array(self.dict[self.beg_str])
+
+        if beg_or_end == "end":
+            return np.array(self.dict[self.end_str])
+
+        raise TypeError(f"unsupport type {beg_or_end} in get_beg_end_flag_idx")
 
     def add_special_char(self, dict_character):
         self.beg_str = "sos"