Upload 4 files

.gitattributes

@@ -1,35 +1,38 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,14 +1,14 @@
----
-title: RapidTableDetection
-emoji: 🏢
-colorFrom: indigo
-colorTo: blue
-sdk: gradio
-sdk_version: 5.6.0
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: extract table form complex bg
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: RapidTableDetection
+emoji: 🏢
+colorFrom: indigo
+colorTo: blue
+sdk: gradio
+sdk_version: 5.6.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+short_description: extract table form complex bg
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,182 @@
+import gradio as gr
+import os
+import cv2
+from rapid_table_det.inference import TableDetector
+from rapid_table_det.utils.visuallize import img_loader, visuallize, extract_table_img
+
+example_images = [
+    "images/doc1.png",
+    "images/doc2.jpg",
+    "images/doc3.jpg",
+    "images/doc4.jpg",
+    "images/doc5.jpg",
+    "images/real1.jpg",
+    "images/real2.jpeg",
+    "images/real3.jpg",
+    "images/real4.jpg",
+    "images/real5.jpg"
+]
+
+# 定义模型类型选项
+model_type_options = {
+    "YOLO 目标检测": ["yolo_obj_det"],
+    "Paddle 目标检测": ["paddle_obj_det"],
+    "Paddle 目标检测 (量化)": ["paddle_obj_det_s"],
+    "YOLO 语义分割": ["yolo_edge_det"],
+    "YOLO 语义分割 (小型)": ["yolo_edge_det_s"],
+    "Paddle 语义分割": ["paddle_edge_det"],
+    "Paddle 语义分割 (量化)": ["paddle_edge_det_s"],
+    "Paddle 方向分类": ["paddle_cls_det"]
+}
+
+# 预生成所有可能的 TableDetector 实例
+preinitialized_detectors = {}
+
+for obj_model_type in model_type_options["YOLO 目标检测"] + model_type_options["Paddle 目标检测"] + model_type_options[
+    "Paddle 目标检测 (量化)"]:
+    for edge_model_type in model_type_options["YOLO 语义分割"] + model_type_options["YOLO 语义分割 (小型)"] + model_type_options[
+        "Paddle 语义分割"] + model_type_options["Paddle 语义分割 (量化)"]:
+        for cls_model_type in model_type_options["Paddle 方向分类"]:
+            detector_key = (obj_model_type, edge_model_type, cls_model_type)
+            preinitialized_detectors[detector_key] = TableDetector(
+                obj_model_type=obj_model_type,
+                edge_model_type=edge_model_type,
+                cls_model_type=cls_model_type,
+                obj_model_path=os.path.join("models", f"{obj_model_type}.onnx"),
+                edge_model_path=os.path.join("models", f"{edge_model_type}.onnx"),
+                cls_model_path=os.path.join("models", f"{cls_model_type}.onnx")
+            )
+
+
+# 定义图片缩放函数
+def resize_image(image, max_size=640):
+    height, width = image.shape[:2]
+    if max(height, width) > max_size:
+        scale = max_size / max(height, width)
+        new_height = int(height * scale)
+        new_width = int(width * scale)
+        image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
+    return image
+
+
+# 定义推理函数
+def run_inference(img_path, obj_model_type, edge_model_type, cls_model_type, det_accuracy, use_obj_det, use_edge_det,
+                  use_cls_det):
+    detector_key = (obj_model_type, edge_model_type, cls_model_type)
+    table_det = preinitialized_detectors[detector_key]
+    result, elapse = table_det(
+        img_path,
+        det_accuracy=det_accuracy,
+        use_obj_det=use_obj_det,
+        use_edge_det=use_edge_det,
+        use_cls_det=use_cls_det
+    )
+
+    # 加载图片
+    img = img_loader(img_path)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    extract_img = img.copy()
+    visual_img = img.copy()
+    extract_imgs = []
+
+    for i, res in enumerate(result):
+        box = res["box"]
+        lt, rt, rb, lb = res["lt"], res["rt"], res["rb"], res["lb"]
+        # 带识别框和左上角方向位置
+        visual_img = visuallize(visual_img, box, lt, rt, rb, lb)
+        # 透视变换提取表格图片
+        wrapped_img = extract_table_img(extract_img.copy(), lt, rt, rb, lb)
+        extract_imgs.append(wrapped_img)
+    # 缩放图片
+    visual_img = resize_image(visual_img)
+    extract_imgs = [resize_image(img) for img in extract_imgs]
+
+    obj_det_elapse, edge_elapse, rotate_det_elapse = elapse
+    return visual_img, extract_imgs, f"obj_det_elapse:{obj_det_elapse}, edge_elapse={edge_elapse}, rotate_det_elapse={rotate_det_elapse}"
+
+
+def update_extract_outputs(visual_img, extract_imgs, time_info):
+    if len(extract_imgs) == 1:
+        return visual_img, extract_imgs[0], time_info
+    else:
+        return visual_img, extract_imgs, time_info
+
+
+# 创建Gradio界面
+with gr.Blocks(
+        css="""
+        .scrollable-container {
+            overflow-x: auto;
+            white-space: nowrap;
+        }
+        .header-links {
+            text-align: center;
+        }
+        .header-links a {
+            display: inline-block;
+            text-align: center;
+            margin-right: 10px;  /* 调整间距 */
+        }
+    """
+) as demo:
+    gr.HTML(
+        "<h1 style='text-align: center;'><a href='https://github.com/RapidAI/RapidTableDetection'>RapidTableDetection</a></h1>"
+    )
+    gr.HTML('''
+                                            <div class="header-links">
+                                              <a href=""><img src="https://img.shields.io/badge/Python->=3.8,<3.12-aff.svg"></a>
+  <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Mac%2C%20Win-pink.svg"></a>
+<a href="https://semver.org/"><img alt="SemVer2.0" src="https://img.shields.io/badge/SemVer-2.0-brightgreen"></a>
+  <a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>
+  <a href="https://github.com/RapidAI/TableStructureRec/blob/c41bbd23898cb27a957ed962b0ffee3c74dfeff1/LICENSE"><img alt="GitHub" src="https://img.shields.io/badge/license-Apache 2.0-blue"></a>
+                                            </div>
+                                            ''')
+    with gr.Row():
+        with gr.Column(variant="panel", scale=1):
+            img_input = gr.Image(label="Upload or Select Image", sources="upload", value="images/real1.jpg")
+
+            # 示例图片选择器
+            examples = gr.Examples(
+                examples=example_images,
+                examples_per_page=len(example_images),
+                inputs=img_input,
+                fn=lambda x: x,  # 简单返回图片路径
+                outputs=img_input,
+                cache_examples=False
+            )
+
+            obj_model_type = gr.Dropdown(
+                choices=model_type_options["YOLO 目标检测"] + model_type_options["Paddle 目标检测"] +
+                        model_type_options["Paddle 目标检测 (量化)"],
+                value="yolo_obj_det",
+                label="obj det model")
+            edge_model_type = gr.Dropdown(
+                choices=model_type_options["YOLO 语义分割"] + model_type_options["YOLO 语义分割 (小型)"] +
+                        model_type_options["Paddle 语义分割"] + model_type_options["Paddle 语义分割 (量化)"],
+                value="yolo_edge_det",
+                label="edge seg model")
+            cls_model_type = gr.Dropdown(choices=model_type_options["Paddle 方向分类"],
+                                         value="paddle_cls_det",
+                                         label="direction cls model")
+
+            det_accuracy = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="目标检测置信度阈值")
+            use_obj_det = gr.Checkbox(value=True, label="use obj det")
+            use_edge_det = gr.Checkbox(value=True, label="use edge seg")
+            use_cls_det = gr.Checkbox(value=True, label="use direction cls")
+
+            run_button = gr.Button("run")
+
+        with gr.Column(scale=2):
+            visual_output = gr.Image(label="output visualize")
+            extract_outputs = gr.Gallery(label="extracted images", object_fit="contain", columns=1, preview=True)
+            time_output = gr.Textbox(label="elapsed")
+
+    run_button.click(
+        fn=run_inference,
+        inputs=[img_input, obj_model_type, edge_model_type, cls_model_type, det_accuracy, use_obj_det, use_edge_det,
+                use_cls_det],
+        outputs=[visual_output, extract_outputs, time_output]
+    )
+
+# 启动Gradio应用
+demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+gradio
+rapid-table-det