fix: 修正偶发异常

demo_onnx.py

@@ -5,14 +5,15 @@
 from rapid_table_det.inference import TableDetector
 from rapid_table_det.utils import visuallize, extract_table_img, img_loader
 
-img_path = f"images/WechatIMG148.jpeg"
+img_path = f"images/web1.png"
 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)
 table_det = TableDetector(
-    obj_model_path="rapid_table_det/models/obj_det.onnx",
+    obj_model_path="rapid_table_det/models/modified_obj_det.onnx",
+    # obj_model_path="rapid_table_det/models/obj_det.onnx",
     edge_model_path="rapid_table_det/models/edge_det.onnx",
     cls_model_path="rapid_table_det/models/cls_det.onnx",
     use_obj_det=True,
@@ -26,12 +27,13 @@
 )
 # 一张图片中可能有多个表格
 img = img_loader(img_path)
+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)
     # 透视变换提取表格图片
-    extract_img = extract_table_img(img.copy(), lt, rt, rb, lb)
+    extract_img = extract_table_img(extract_img.copy(), lt, rt, rb, lb)
     cv2.imwrite(f"{out_dir}/{file_name}-extract-{i}.jpg", extract_img)
 cv2.imwrite(f"{out_dir}/{file_name}-visualize.jpg", img)

demo_paddle.py

@@ -5,16 +5,16 @@
 from rapid_table_det_paddle.inference import TableDetector
 from rapid_table_det_paddle.utils import visuallize, extract_table_img, img_loader
 
-img_path = f"images/car-invoice-img02282.jpg"
+img_path = f"images/wtw_poc-00013.jpg"
 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)
 table_det = TableDetector(
-    obj_model_path="models/obj_det/model",
-    edge_model_path="models/db_net/model",
-    cls_model_path="models/pplcnet/model",
+    obj_model_path="rapid_table_det_paddle/models/obj_det/model",
+    edge_model_path="rapid_table_det_paddle/models/db_net/model",
+    cls_model_path="rapid_table_det_paddle/models/pplcnet/model",
     use_obj_det=True,
     use_edge_det=True,
     use_rotate_det=True,
@@ -26,12 +26,13 @@
 )
 # 一张图片中可能有多个表格
 img = img_loader(img_path)
+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)
     # 透视变换提取表格图片
-    extract_img = extract_table_img(img.copy(), lt, rt, rb, lb)
+    extract_img = extract_table_img(extract_img.copy(), lt, rt, rb, lb)
     cv2.imwrite(f"{out_dir}/{file_name}-extract-{i}.jpg", extract_img)
 cv2.imwrite(f"{out_dir}/{file_name}-visualize.jpg", img)

rapid_table_det/utils.py

@@ -419,8 +419,6 @@ def get_max_adjacent_bbox(mask):
 
 
 def visuallize(img, box, lt, rt, rb, lb):
-    # img = cv2.imread(img_path)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     xmin, ymin, xmax, ymax = box
     draw_box = np.array([lt, rt, rb, lb]).reshape([-1, 2])
     cv2.circle(img, (int(lt[0]), int(lt[1])), 50, (255, 0, 0), 10)
@@ -449,7 +447,6 @@ def extract_table_img(img, lt, rt, rb, lb):
     返回:
     numpy.ndarray: 提取出的角点区域图片
     """
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     # 源点坐标
     src_points = np.float32([lt, rt, lb, rb])
 

rapid_table_det_paddle/predictor.py

@@ -39,7 +39,7 @@ def __call__(self, img, **kwargs):
         return result, time.time() - start
 
     def img_preprocess(self, img, resize_shape=[928, 928]):
-        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         im_info = {
             'scale_factor': np.array(
                 [1., 1.], dtype=np.float32),
@@ -147,7 +147,7 @@ def sort_and_clip_coordinates(self, box):
         return lt, lb, rt, rb
 
     def img_preprocess(self, img, resize_shape=[800, 800]):
-        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         im, new_w, new_h, left, top = ResizePad(img, resize_shape[0])
         im = im / 255.0
         im = im.transpose((2, 0, 1)).copy()
@@ -165,7 +165,7 @@ def __init__(self, model_path=pplcnet_model_path, **kwargs):
     def __call__(self, img, **kwargs):
         start = time.time()
         img = self.img_loader(img)
-        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         img = self.img_preprocess(img, self.resize_shape)
         with paddle.no_grad():
             # print(cls_img.shape)

rapid_table_det_paddle/utils.py

@@ -347,8 +347,6 @@ def get_max_adjacent_bbox(mask):
 
 
 def visuallize(img, box, lt, rt, rb, lb):
-    # img = cv2.imread(img_path)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     xmin, ymin, xmax, ymax = box
     draw_box = np.array([lt, rt, rb, lb]).reshape([-1, 2])
     cv2.circle(img, (int(lt[0]), int(lt[1])), 50, (255, 0, 0), 10)
@@ -377,7 +375,6 @@ def extract_table_img(img, lt, rt, rb, lb):
     返回:
     numpy.ndarray: 提取出的角点区域图片
     """
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     # 源点坐标
     src_points = np.float32([lt, rt, lb, rb])
 
@@ -404,5 +401,4 @@ def extract_table_img(img, lt, rt, rb, lb):
 
     # 应用透视变换
     warped = cv2.warpPerspective(img, M, (max_width, max_height))
-
     return warped