feat: add yolo model and auto download

Author: Joker1212

README.md

@@ -10,63 +10,82 @@
 </div>
 
 ### 最近更新
+
 - **2024.10.15**
     - 完成初版代码，包含目标检测，语义分割，角点方向识别三个模块
+- **2024.11.2**
+    - 补充新训练yolo11的目标检测模型和边缘检测模型，增加自动下载，轻量化包体积，自由组合各个模块
 
 ### 简介
+
 💡✨ 强大且高效的表格检测，支持论文、期刊、杂志、发票、收据、签到单等各种表格。
 
-🚀 支持高精度 Paddle 版本和量化 ONNX 版本，单图 CPU 推理仅需 1.5 秒，Paddle-GPU(V100) 仅需 0.2 秒。
+🚀 支持来源于paddle和yolo的版本，平衡速度和精度下单图 CPU 推理仅需 1 秒，Paddle-GPU(V100) 仅需 0.2 秒。
 
 🛠️ 支持三个模块自由组合，独立训练调优，提供 ONNX 转换脚本和微调训练方案。
 
 🌟 whl 包轻松集成使用，为下游 OCR、表格识别和数据采集提供强力支撑。
 
-📚参考项目 [百度表格检测大赛第2名方案](https://aistudio.baidu.com/projectdetail/5398861?searchKeyword=%E8%A1%A8%E6%A0%BC%E6%A3%80%E6%B5%8B%E5%A4%A7%E8%B5%9B&searchTab=ALL) 的实现方案，补充大量真实场景数据再训练
+📚参考项目 [百度表格检测大赛第2名方案](https://aistudio.baidu.com/projectdetail/5398861?searchKeyword=%E8%A1%A8%E6%A0%BC%E6%A3%80%E6%B5%8B%E5%A4%A7%E8%B5%9B&searchTab=ALL)
+的实现方案，补充大量真实场景数据再训练
 ![img.png](readme_resource/structure.png)
 👇🏻训练数据集在致谢, 作者天天上班摸鱼搞开源，希望大家点个⭐️支持一下
 
 ### 使用建议
+
 📚 文档场景: 无透视旋转，只使用目标检测\
 📷 拍照场景小角度旋转(-90~90): 默认左上角，不使用角点方向识别\
 🔍 使用在线体验找到适合你场景的模型组合
-### 在线体验
 
+### 在线体验
 
 ### 效果展示
+
 ![res_show.jpg](readme_resource/res_show.jpg)![res_show2.jpg](readme_resource/res_show2.jpg)
+
 ### 安装
-为简化使用，已经将最小的量化模型打包到 rapid_table_det 中，需要更高精度或gpu推理，请自行下载对应模型
-🪜下载模型 [modescope模型仓](https://www.modelscope.cn/models/jockerK/TableExtractor) [release assets](https://github.com/Joker1212/RapidTableDetection/releases/tag/v0.0.0)
+
+🪜模型会自动下载，也可以自己去仓库下载 [modescope模型仓](https://www.modelscope.cn/models/jockerK/TableExtractor)
+
 ``` python {linenos=table}
 # 建议使用清华源安装 https://pypi.tuna.tsinghua.edu.cn/simple
 pip install rapid-table-det
 ```
 
 #### 参数说明
-cpu和gpu的初始化完全一致
-```
-table_det = TableDetector(
-# 目标检测表格模型
-obj_model_path="models/obj_det_paddle(obj_det.onnx)", 
-# 边角检测表格模型(从复杂环境得到表格多边形框)
-edge_model_path="models/edge_det_paddle(edge_det.onnx)", 
-# 角点方向识别
-cls_model_path="models/cls_det_paddle(cls_det.onnx)", 
-# 文档场景已经由版面识别模型提取设置为False
-use_obj_det=True, 
-# 只有90，180，270大角度旋转且无透视时候设置为False
-use_edge_det=True, 
-# 小角度(-90~90)旋转设置为False
-use_cls_det=True, 
-)
-```
+
+默认值
+use_cuda: False : 启用gpu加速推理 \
+obj_model_type="yolo_obj_det", \
+edge_model_type= "yolo_edge_det", \
+cls_model_type= "paddle_cls_det"
+
+| `model_type`         | 任务类型   | 训练来源                                 | 大小     | 单表格耗时                 |
+|:---------------------|:-------|:-------------------------------------|:-------|:----------------------|
+| **yolo_obj_det**     | 表格目标检测 | `yolo11-l`                           | `100m` | `cpu:500ms, gpu:0.2`  |
+| `paddle_obj_det`     | 表格目标检测 | `paddle yoloe-plus-x`                | `380m` | `cpu:500ms, gpu:0.2`  |
+| `paddle_obj_det_s`   | 表格目标检测 | `paddle yoloe-plus-x + quantization` | `95m`  | `cpu:1000ms, gpu:0.2` |
+| **yolo_edge_det**    | 语义分割   | `yolo11-l-segment`                   | `108m` | `cpu:500ms, gpu:0.2`  |
+| `yolo_edge_det_s`    | 语义分割   | `yolo11-s-segment`                   | `11m`  | `cpu:100ms, gpu:0.2`  |
+| `paddle_edge_det`    | 语义分割   | `paddle-dbnet`                       | `99m`  | `cpu:600ms, gpu:0.2`  |
+| `paddle_edge_det_s`  | 语义分割   | `paddle-dbnet + quantization`        | `25m`  | `cpu:500ms, gpu:0.2`  |
+| **paddle_cls_det**     | 方向分类   | `paddle pplcnet`                     | `6.5m` | `cpu:70ms, gpu:0.2`   |
+
+
+执行参数
+det_accuracy=0.7,
+use_obj_det=True,
+use_edge_det=True,
+use_cls_det=True,
 
 ### 快速使用
+
 ``` python {linenos=table}
 from rapid_table_det.inference import TableDetector
+
+img_path = f"images/weixin.png"
 table_det = TableDetector()
-img_path = f"tests/test_files/chip.jpg"
+
 result, elapse = table_det(img_path)
 obj_det_elapse, edge_elapse, rotate_det_elapse = elapse
 print(
@@ -75,7 +94,8 @@ print(
 # 输出可视化
 # import os
 # import cv2
-# from rapid_table_det.utils import img_loader, visuallize, extract_table_img
+# from rapid_table_det.utils.visuallize import img_loader, visuallize, extract_table_img
+# 
 # img = img_loader(img_path)
 # file_name_with_ext = os.path.basename(img_path)
 # file_name, file_ext = os.path.splitext(file_name_with_ext)
@@ -94,68 +114,28 @@ print(
 # cv2.imwrite(f"{out_dir}/{file_name}-visualize.jpg", img)
 
 ```
-### gpu版本使用
-必须下载模型，指定模型位置！
-``` python {linenos=table}
-# 建议使用清华源安装 https://pypi.tuna.tsinghua.edu.cn/simple
-pip install rapid-table-det-paddle (默认安装gpu版本，可以自行覆盖安装cpu版本paddlepaddle)
-```
-```python
-from rapid_table_det_paddle.inference import TableDetector
-
-img_path = f"tests/test_files/chip.jpg"
-
-table_det = TableDetector(
-    obj_model_path="models/obj_det_paddle",
-    edge_model_path="models/edge_det_paddle",
-    cls_model_path="models/cls_det_paddle",
-    use_obj_det=True,
-    use_edge_det=True,
-    use_cls_det=True,
-)
-result, elapse = table_det(img_path)
-obj_det_elapse, edge_elapse, rotate_det_elapse = elapse
-print(
-    f"obj_det_elapse:{obj_det_elapse}, edge_elapse={edge_elapse}, rotate_det_elapse={rotate_det_elapse}"
-)
-# 一张图片中可能有多个表格
-# img = img_loader(img_path)
-# file_name_with_ext = os.path.basename(img_path)
-# file_name, file_ext = os.path.splitext(file_name_with_ext)
-# out_dir = "rapid_table_det_paddle/outputs"
-# if not os.path.exists(out_dir):
-#     os.makedirs(out_dir)
-# extract_img = img.copy()
-# for i, res in enumerate(result):
-#     box = res["box"]
-#     lt, rt, rb, lb = res["lt"], res["rt"], res["rb"], res["lb"]
-#     # 带识别框和左上角方向位置
-#     img = visuallize(img, box, lt, rt, rb, lb)
-#     # 透视变换提取表格图片
-#     wrapped_img = extract_table_img(extract_img.copy(), lt, rt, rb, lb)
-#     cv2.imwrite(f"{out_dir}/{file_name}-extract-{i}.jpg", wrapped_img)
-# cv2.imwrite(f"{out_dir}/{file_name}-visualize.jpg", img)
-
-```
-
 
 ## FAQ (Frequently Asked Questions)
 
 1. **问：如何微调模型适应特定场景?**
-    - 答：直接参考这个项目，有非常详细的可视化操作步骤,可以得到paddle的推理模型 [百度表格检测大赛](https://aistudio.baidu.com/projectdetail/5398861?searchKeyword=%E8%A1%A8%E6%A0%BC%E6%A3%80%E6%B5%8B%E5%A4%A7%E8%B5%9B&searchTab=ALL) 
+   -
+   答：直接参考这个项目，有非常详细的可视化操作步骤,可以得到paddle的推理模型 [百度表格检测大赛](https://aistudio.baidu.com/projectdetail/5398861?searchKeyword=%E8%A1%A8%E6%A0%BC%E6%A3%80%E6%B5%8B%E5%A4%A7%E8%B5%9B&searchTab=ALL)
 
 2. **问：如何导出onnx**
-   - 答：在本项目tools下，有onnx_transform.ipynb文件，可以照步骤执行(因为pp-yoloe导出onnx有bug一直没修，这里我自己写了一个fix_onnx脚本改动onnx模型节点来临时解决了)
+    - 答：在本项目tools下，有onnx_transform.ipynb文件，可以照步骤执行(
+      因为pp-yoloe导出onnx有bug一直没修，这里我自己写了一个fix_onnx脚本改动onnx模型节点来临时解决了)
 
 3. **问：图片有扭曲可以修正吗？**
     - 答：本项目只解决旋转和透视场景的表格提取，对于扭曲的场景，需要先进行扭曲修正
 
 ### 致谢
+
 [百度表格检测大赛第2名方案](https://aistudio.baidu.com/projectdetail/5398861?searchKeyword=%E8%A1%A8%E6%A0%BC%E6%A3%80%E6%B5%8B%E5%A4%A7%E8%B5%9B&searchTab=ALL) \
 [WTW 自然场景表格数据集](https://tianchi.aliyun.com/dataset/108587) \
 [FinTabNet PDF文档表格数据集](https://developer.ibm.com/exchanges/data/all/fintabnet/) \
 [TableBank 表格数据集](https://doc-analysis.github.io/tablebank-page/) \
 [TableGeneration 表格自动生成工具](https://github.com/WenmuZhou/TableGeneration)
+
 ### 贡献指南
 
 欢迎提交请求。对于重大更改，请先打开issue讨论您想要改变的内容。

demo_onnx.py

@@ -1,35 +1,33 @@
 from rapid_table_det.inference import TableDetector
 
-# img_path = f"tests/test_files/chip2.jpg"
 img_path = f"images/weixin.png"
 table_det = TableDetector(
-    obj_model_path="rapid_table_det/models/yolo_obj_det_l.onnx",
-    edge_model_path="rapid_table_det/models/yolo_edge_det_s.onnx",
+    obj_model_type="paddle_obj_det_s", edge_model_type="paddle_edge_det_s"
 )
-if __name__ == "__main__":
-    result, elapse = table_det(img_path)
-    obj_det_elapse, edge_elapse, rotate_det_elapse = elapse
-    print(
-        f"obj_det_elapse:{obj_det_elapse}, edge_elapse={edge_elapse}, rotate_det_elapse={rotate_det_elapse}"
-    )
-    # 输出可视化
-    import os
-    import cv2
-    from rapid_table_det.utils import img_loader, visuallize, extract_table_img
 
-    img = img_loader(img_path)
-    file_name_with_ext = os.path.basename(img_path)
-    file_name, file_ext = os.path.splitext(file_name_with_ext)
-    out_dir = "rapid_table_det/outputs"
-    if not os.path.exists(out_dir):
-        os.makedirs(out_dir)
-    extract_img = img.copy()
-    for i, res in enumerate(result):
-        box = res["box"]
-        lt, rt, rb, lb = res["lt"], res["rt"], res["rb"], res["lb"]
-        # 带识别框和左上角方向位置
-        img = visuallize(img, box, lt, rt, rb, lb)
-        # 透视变换提取表格图片
-        wrapped_img = extract_table_img(extract_img.copy(), lt, rt, rb, lb)
-        cv2.imwrite(f"{out_dir}/{file_name}-extract-{i}.jpg", wrapped_img)
-    cv2.imwrite(f"{out_dir}/{file_name}-visualize.jpg", img)
+result, elapse = table_det(img_path)
+obj_det_elapse, edge_elapse, rotate_det_elapse = elapse
+print(
+    f"obj_det_elapse:{obj_det_elapse}, edge_elapse={edge_elapse}, rotate_det_elapse={rotate_det_elapse}"
+)
+# 输出可视化
+# import os
+# import cv2
+# from rapid_table_det.utils.visuallize import img_loader, visuallize, extract_table_img
+#
+# img = img_loader(img_path)
+# file_name_with_ext = os.path.basename(img_path)
+# file_name, file_ext = os.path.splitext(file_name_with_ext)
+# out_dir = "rapid_table_det/outputs"
+# if not os.path.exists(out_dir):
+#     os.makedirs(out_dir)
+# extract_img = img.copy()
+# for i, res in enumerate(result):
+#     box = res["box"]
+#     lt, rt, rb, lb = res["lt"], res["rt"], res["rb"], res["lb"]
+#     # 带识别框和左上角方向位置
+#     img = visuallize(img, box, lt, rt, rb, lb)
+#     # 透视变换提取表格图片
+#     wrapped_img = extract_table_img(extract_img.copy(), lt, rt, rb, lb)
+#     cv2.imwrite(f"{out_dir}/{file_name}-extract-{i}.jpg", wrapped_img)
+# cv2.imwrite(f"{out_dir}/{file_name}-visualize.jpg", img)

rapid_table_det/__init__.py

@@ -2,7 +2,3 @@
 # @Author: Jocker1212
 # @Contact: xinyijianggo@gmail.com
 from .inference import TableDetector
-from .utils import img_loader, visuallize, extract_table_img
-
-#
-__all__ = ["TableDetector", "img_loader", "visuallize", "extract_table_img"]

rapid_table_det/inference.py

@@ -1,35 +1,75 @@
 import os
 from pathlib import Path
+from typing import Union
 
 import cv2
 import numpy as np
 
-from rapid_table_det.predictor import DbNet, ObjectDetector, PPLCNet, YoloSeg, YoloDet
-from rapid_table_det.utils import LoadImage
+from .predictor import DbNet, PaddleYoloEDet, PPLCNet, YoloSeg, YoloDet
+from .utils.download_model import DownloadModel
+
+from .utils.logger import get_logger
+from .utils.load_image import LoadImage
 
 root_dir = Path(__file__).resolve().parent
 model_dir = os.path.join(root_dir, "models")
 
+ROOT_DIR = Path(__file__).resolve().parent
+logger = get_logger("rapid_layout")
+
+ROOT_URL = "https://www.modelscope.cn/models/jockerK/TableExtractor/resolve/master/rapid_table_det/models/"
+KEY_TO_MODEL_URL = {
+    "yolo_obj_det": f"{ROOT_URL}/yolo_obj_det.onnx",
+    "yolo_edge_det": f"{ROOT_URL}/yolo_edge_det.onnx",
+    "yolo_edge_det_s": f"{ROOT_URL}/yolo_edge_det_s.onnx",
+    "paddle_obj_det": f"{ROOT_URL}/paddle_obj_det.onnx",
+    "paddle_obj_det_s": f"{ROOT_URL}/paddle_obj_det_s.onnx",
+    "paddle_edge_det": f"{ROOT_URL}/paddle_edge_det.onnx",
+    "paddle_edge_det_s": f"{ROOT_URL}/paddle_edge_det_s.onnx",
+    "paddle_cls_det": f"{ROOT_URL}/paddle_cls_det.onnx",
+}
+
 
 class TableDetector:
     def __init__(
         self,
-        obj_model="yolo",
-        edge_model="yolo",
-        obj_model_path=os.path.join(model_dir, "obj_det_quantized.onnx"),
-        edge_model_path=os.path.join(model_dir, "yolo_edge_det_s.onnx"),
-        cls_model_path=os.path.join(model_dir, "cls_det.onnx"),
+        use_cuda=False,
+        use_dml=False,
+        obj_model_path=None,
+        edge_model_path=None,
+        cls_model_path=None,
+        obj_model_type="yolo_obj_det",
+        edge_model_type="yolo_edge_det",
+        cls_model_type="paddle_cls_det",
     ):
         self.img_loader = LoadImage()
-        if obj_model == "yolo":
-            self.obj_detector = YoloDet(obj_model_path)
+        obj_det_config = {
+            "model_path": self.get_model_path(obj_model_type, obj_model_path),
+            "use_cuda": use_cuda,
+            "use_dml": use_dml,
+        }
+        edge_det_config = {
+            "model_path": self.get_model_path(edge_model_type, edge_model_path),
+            "use_cuda": use_cuda,
+            "use_dml": use_dml,
+        }
+        cls_det_config = {
+            "model_path": self.get_model_path(cls_model_type, cls_model_path),
+            "use_cuda": use_cuda,
+            "use_dml": use_dml,
+        }
+        if "yolo" in obj_model_type:
+            self.obj_detector = YoloDet(obj_det_config)
         else:
-            self.obj_detector = ObjectDetector(obj_model_path)
-        if edge_model == "yolo":
-            self.dbnet = YoloSeg(edge_model_path)
+            self.obj_detector = PaddleYoloEDet(obj_det_config)
+        if "yolo" in edge_model_type:
+            self.dbnet = YoloSeg(edge_det_config)
         else:
-            self.dbnet = DbNet(edge_model_path)
-        self.pplcnet = PPLCNet(cls_model_path)
+            self.dbnet = DbNet(edge_det_config)
+        if "yolo" in cls_model_type:
+            self.pplcnet = PPLCNet(cls_det_config)
+        else:
+            self.pplcnet = PPLCNet(cls_det_config)
 
     def __call__(
         self,
@@ -101,6 +141,16 @@ def init_default_output(self, h, w):
         return obj_det_res, pred_label
 
     def add_pre_info_for_cls(self, cls_img, edge_box, xmin_cls, ymin_cls):
+        """
+        Args:
+            cls_img:
+            edge_box:
+            xmin_cls:
+            ymin_cls:
+
+        Returns: 带边缘划线的图片，给方向分类提供先验信息
+
+        """
         cls_box = edge_box.copy()
         cls_box[:, 0] = cls_box[:, 0] - xmin_cls
         cls_box[:, 1] = cls_box[:, 1] - ymin_cls
@@ -166,3 +216,18 @@ def pad_box_points(self, h, w, xmax, xmin, ymax, ymin, pad):
         ymax_edge = min(ymax + pad, h)
         xmax_edge = min(xmax + pad, w)
         return xmin_edge, ymin_edge, xmax_edge, ymax_edge
+
+    @staticmethod
+    def get_model_path(model_type: str, model_path: Union[str, Path, None]) -> str:
+        if model_path is not None:
+            return model_path
+
+        model_url = KEY_TO_MODEL_URL.get(model_type, None)
+        if model_url:
+            model_path = DownloadModel.download(model_url)
+            return model_path
+
+        logger.info(
+            "model url is None, using the default download model %s", model_path
+        )
+        return model_path