feat: add yolo model

Author: Joker1212

demo_onnx.py

@@ -1,33 +1,35 @@
 from rapid_table_det.inference import TableDetector
 
-img_path = f"tests/test_files/chip2.jpg"
+# img_path = f"tests/test_files/chip2.jpg"
+img_path = f"images/weixin.png"
 table_det = TableDetector(
-    obj_model_path="rapid_table_det/models/obj_det.onnx",
-    edge_model_path="rapid_table_det/models/edge_det.onnx",
+    obj_model_path="rapid_table_det/models/yolo_obj_det_l.onnx",
+    edge_model_path="rapid_table_det/models/yolo_edge_det_s.onnx",
 )
-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
+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)
+    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/inference.py

@@ -4,71 +4,57 @@
 import cv2
 import numpy as np
 
-from rapid_table_det.predictor import DbNet, ObjectDetector, PPLCNet
+from rapid_table_det.predictor import DbNet, ObjectDetector, PPLCNet, YoloSeg, YoloDet
 from rapid_table_det.utils import LoadImage
 
-MODEL_URLS = {
-    "onnx_tiny": {
-        "obj_det": "https://github.com/Joker1212/RapidTableDetection/releases/download/v0.0.0/obj_det_quantized.zip",
-        "edge_det": "https://github.com/Joker1212/RapidTableDetection/releases/download/v0.0.0/edge_det_quantized.zip",
-        "cls_det": "https://github.com/Joker1212/RapidTableDetection/releases/download/v0.0.0/cls_det.zip",
-    },
-    "onnx": {
-        "obj_det": "https://github.com/Joker1212/RapidTableDetection/releases/download/v0.0.0/obj_det.zip",
-        "edge_det": "https://github.com/Joker1212/RapidTableDetection/releases/download/v0.0.0/edge_det.zip",
-        "cls_det": "https://github.com/Joker1212/RapidTableDetection/releases/download/v0.0.0/cls_det.zip",
-    },
-}
 root_dir = Path(__file__).resolve().parent
 model_dir = os.path.join(root_dir, "models")
 
 
 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, "edge_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_obj_det=True,
-        use_edge_det=True,
-        use_cls_det=True,
     ):
-        self.use_obj_det = use_obj_det
-        self.use_edge_det = use_edge_det
-        self.use_cls_det = use_cls_det
         self.img_loader = LoadImage()
-
-        if self.use_obj_det:
+        if obj_model == "yolo":
+            self.obj_detector = YoloDet(obj_model_path)
+        else:
             self.obj_detector = ObjectDetector(obj_model_path)
-        if self.use_edge_det:
+        if edge_model == "yolo":
+            self.dbnet = YoloSeg(edge_model_path)
+        else:
             self.dbnet = DbNet(edge_model_path)
-        if self.use_cls_det:
-            self.pplcnet = PPLCNet(cls_model_path)
+        self.pplcnet = PPLCNet(cls_model_path)
 
-    def __call__(self, img, det_accuracy=0.7):
+    def __call__(
+        self,
+        img,
+        det_accuracy=0.7,
+        use_obj_det=True,
+        use_edge_det=True,
+        use_cls_det=True,
+    ):
         img = self.img_loader(img)
         img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         img_mask = img.copy()
         h, w = img.shape[:-1]
-        img_box = np.array([0, 0, w, h])
-        lb, lt, rb, rt = self.get_box_points(img_box)
-        # 初始化默认值
-        obj_det_res, edge_box, pred_label = (
-            [[1.0, img_box]],
-            img_box.reshape([-1, 2]),
-            0,
-        )
+        obj_det_res, pred_label = self.init_default_output(h, w)
         result = []
         obj_det_elapse, edge_elapse, rotate_det_elapse = 0, 0, 0
-        if self.use_obj_det:
+        if use_obj_det:
             obj_det_res, obj_det_elapse = self.obj_detector(img, score=det_accuracy)
         for i in range(len(obj_det_res)):
             det_res = obj_det_res[i]
             score, box = det_res
             xmin, ymin, xmax, ymax = box
             edge_box = box.reshape([-1, 2])
             lb, lt, rb, rt = self.get_box_points(box)
-            if self.use_edge_det:
+            if use_edge_det:
                 xmin_edge, ymin_edge, xmax_edge, ymax_edge = self.pad_box_points(
                     h, w, xmax, xmin, ymax, ymin, 10
                 )
@@ -77,30 +63,16 @@ def __call__(self, img, det_accuracy=0.7):
                 edge_elapse += tmp_edge_elapse
                 if edge_box is None:
                     continue
-                edge_box[:, 0] += xmin_edge
-                edge_box[:, 1] += ymin_edge
-                lt, lb, rt, rb = (
-                    lt + [xmin_edge, ymin_edge],
-                    lb + [xmin_edge, ymin_edge],
-                    rt + [xmin_edge, ymin_edge],
-                    rb + [xmin_edge, ymin_edge],
+                lb, lt, rb, rt = self.adjust_edge_points_axis(
+                    edge_box, lb, lt, rb, rt, xmin_edge, ymin_edge
                 )
-            if self.use_cls_det:
+            if use_cls_det:
                 xmin_cls, ymin_cls, xmax_cls, ymax_cls = self.pad_box_points(
                     h, w, xmax, xmin, ymax, ymin, 5
                 )
-                cls_box = edge_box.copy()
                 cls_img = img_mask[ymin_cls:ymax_cls, xmin_cls:xmax_cls, :]
-                cls_box[:, 0] = cls_box[:, 0] - xmin_cls
-                cls_box[:, 1] = cls_box[:, 1] - ymin_cls
-                # 画框增加先验信息，辅助方向label识别
-                cv2.polylines(
-                    cls_img,
-                    [np.array(cls_box).astype(np.int32).reshape((-1, 1, 2))],
-                    True,
-                    color=(255, 0, 255),
-                    thickness=5,
-                )
+                # 增加先验信息
+                self.add_pre_info_for_cls(cls_img, edge_box, xmin_cls, ymin_cls)
                 pred_label, tmp_rotate_det_elapse = self.pplcnet(cls_img)
                 rotate_det_elapse += tmp_rotate_det_elapse
             lb1, lt1, rb1, rt1 = self.get_real_rotated_points(
@@ -118,6 +90,40 @@ def __call__(self, img, det_accuracy=0.7):
         elapse = [obj_det_elapse, edge_elapse, rotate_det_elapse]
         return result, elapse
 
+    def init_default_output(self, h, w):
+        img_box = np.array([0, 0, w, h])
+        # 初始化默认值
+        obj_det_res, edge_box, pred_label = (
+            [[1.0, img_box]],
+            img_box.reshape([-1, 2]),
+            0,
+        )
+        return obj_det_res, pred_label
+
+    def add_pre_info_for_cls(self, cls_img, edge_box, xmin_cls, ymin_cls):
+        cls_box = edge_box.copy()
+        cls_box[:, 0] = cls_box[:, 0] - xmin_cls
+        cls_box[:, 1] = cls_box[:, 1] - ymin_cls
+        # 画框增加先验信息，辅助方向label识别
+        cv2.polylines(
+            cls_img,
+            [np.array(cls_box).astype(np.int32).reshape((-1, 1, 2))],
+            True,
+            color=(255, 0, 255),
+            thickness=5,
+        )
+
+    def adjust_edge_points_axis(self, edge_box, lb, lt, rb, rt, xmin_edge, ymin_edge):
+        edge_box[:, 0] += xmin_edge
+        edge_box[:, 1] += ymin_edge
+        lt, lb, rt, rb = (
+            lt + [xmin_edge, ymin_edge],
+            lb + [xmin_edge, ymin_edge],
+            rt + [xmin_edge, ymin_edge],
+            rb + [xmin_edge, ymin_edge],
+        )
+        return lb, lt, rb, rt
+
     def get_box_points(self, img_box):
         x1, y1, x2, y2 = img_box
         lt = np.array([x1, y1])  # 左上角