add table-transformer-detection example

Signed-off-by: yuwenzho <[email protected]>

Author: yuwenzho

examples/.config/model_params_onnxrt.json

@@ -868,10 +868,17 @@
       "main_script": "main.py",
       "batch_size": 1
     },
-    "table_transformer": {
+    "table_transformer_structure_recognition": {
       "model_src_dir": "object_detection/table_transformer/quantization/ptq_static",
-      "dataset_location": "/tf_dataset/dataset/PubTables-1M-Structure",
-      "input_model": "/tf_dataset2/models/onnx/table-transformer/model.onnx",
+      "dataset_location": "/tf_dataset/dataset/PubTables-1M",
+      "input_model": "/tf_dataset2/models/onnx/table-transformer/pubtables1m_structure_detr_r18.onnx",
+      "main_script": "table-transformer/src/main.py",
+      "batch_size": 1
+    },
+    "table_transformer_detection": {
+      "model_src_dir": "object_detection/table_transformer/quantization/ptq_static",
+      "dataset_location": "/tf_dataset/dataset/PubTables-1M",
+      "input_model": "/tf_dataset2/models/onnx/table-transformer/pubtables1m_detection_detr_r18.onnx",
       "main_script": "table-transformer/src/main.py",
       "batch_size": 1
     },

examples/README.md

@@ -1402,7 +1402,13 @@ Intel® Neural Compressor validated examples with multiple compression technique
     <td><a href="./onnxrt/object_detection/ssd_mobilenet_v2/quantization/ptq_static">qlinearops</a> / <a href="./onnxrt/object_detection/ssd_mobilenet_v2/quantization/ptq_static">qdq</a></td>
   </tr>
   <tr>
-    <td>Table Transformer</td>
+    <td>Table Transformer Structure Recognition</td>
+    <td>Object Detection</td>
+    <td>Post-Training Static Quantization</td>
+    <td><a href="./onnxrt/object_detection/table_transformer/quantization/ptq_static">qlinearops</a></td>
+  </tr>
+  <tr>
+    <td>Table Transformer Detection</td>
     <td>Object Detection</td>
     <td>Post-Training Static Quantization</td>
     <td><a href="./onnxrt/object_detection/table_transformer/quantization/ptq_static">qlinearops</a></td>

examples/onnxrt/object_detection/table_transformer/quantization/ptq_static/README.md

@@ -1,7 +1,7 @@
 Step-by-Step
 ============
 
-This example show how to export, quantize and evaluate the DETR R18 model for table structure recognition task based on PubTables-1M dataset.
+This example show how to export, quantize and evaluate 2 [DETR](https://huggingface.co/docs/transformers/model_doc/detr) R18 models on [PubTables-1M](https://huggingface.co/datasets/bsmock/pubtables-1m) dataset, one for table detection and one for table structure recognition, dubbed Table Transformers.
 
 # Prerequisite
 
@@ -16,16 +16,20 @@ bash prepare.sh
 
 ## 2. Prepare Dataset
 
-Download dataset according to this [doc](https://github.com/microsoft/table-transformer/tree/main#training-and-evaluation-data).
+Download PubTables-1M dataset according to this [doc](https://github.com/microsoft/table-transformer/tree/main#training-and-evaluation-data).
+After downloading and extracting, PubTables-1M dataset folder should contain `PubTables-1M-Structure` and `PubTables-1M-Detection` folders.
 
 ## 3. Prepare Model
 
-```shell
-wget https://huggingface.co/bsmock/tatr-pubtables1m-v1.0/resolve/main/pubtables1m_structure_detr_r18.pth
+Prepare DETR R18 model for table structure recognition.
 
-bash export.sh --input_model=/path/to/pubtables1m_structure_detr_r18.pth \
-               --output_model=/path/to/export \ # model path as *.onnx
-               --dataset_location=/path/to/dataset_folder # dataset_folder should contains 'words' sub-folder
+```
+python prepare_model.py  --input_model=structure_detr  --output_model=pubtables1m_structure_detr_r18.onnx --dataset_location=/path/to/pubtables-1m
+```
+
+Prepare DETR R18 model for table detection.
+```
+python prepare_model.py  --input_model=detection_detr  --output_model=pubtables1m_detection_detr_r18.onnx --dataset_location=/path/to/pubtables-1m
 ```
 
 # Run
@@ -37,13 +41,13 @@ Static quantization with QOperator format:
 ```bash
 bash run_tuning.sh --input_model=path/to/model  \ # model path as *.onnx
                    --output_model=path/to/save \ # model path as *.onnx
-                   --dataset_location=/path/to/dataset_folder # dataset_folder should contains 'words' sub-folder
+                   --dataset_location=/path/to/pubtables-1m # dataset_folder should contains `PubTables-1M-Structure` and/or `PubTables-1M-Detection` folders
 ```
 
 ## 2. Benchmark
 
 ```bash
 bash run_benchmark.sh --input_model=path/to/model  \ # model path as *.onnx
-                      --dataset_location=/path/to/dataset_folder # dataset_folder should contains 'words' sub-folder
+                      --dataset_location=/path/to/pubtables-1m # dataset_folder should contains `PubTables-1M-Structure` and/or `PubTables-1M-Detection` folders
                       --mode=performance # or accuracy
 ```

examples/onnxrt/object_detection/table_transformer/quantization/ptq_static/export.sh

@@ -0,0 +1,100 @@
+import argparse
+import os
+import subprocess
+import sys
+from urllib import request
+
+MODEL_URLS = {"structure_detr": "https://huggingface.co/bsmock/tatr-pubtables1m-v1.0/resolve/main/pubtables1m_structure_detr_r18.pth",
+              "detection_detr": "https://huggingface.co/bsmock/tatr-pubtables1m-v1.0/resolve/main/pubtables1m_detection_detr_r18.pth"}
+MAX_TIMES_RETRY_DOWNLOAD = 5
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_model",
+                        type=str,
+                        required=False,
+                        choices=["structure_detr", "detection_detr"],
+                        default="structure_detr")
+    parser.add_argument("--output_model", type=str, required=True)
+    parser.add_argument("--dataset_location", type=str, required=True)
+    return parser.parse_args()
+
+
+def progressbar(cur, total=100):
+    percent = '{:.2%}'.format(cur / total)
+    sys.stdout.write("\r[%-100s] %s" % ('#' * int(cur), percent))
+    sys.stdout.flush()
+
+
+def schedule(blocknum, blocksize, totalsize):
+    if totalsize == 0:
+        percent = 0
+    else:
+        percent = min(1.0, blocknum * blocksize / totalsize) * 100
+    progressbar(percent)
+
+
+def download_model(url, retry_times=5):
+    model_name = url.split("/")[-1]
+    if os.path.isfile(model_name):
+        print(f"{model_name} exists, skip download")
+        return True
+
+    print("download model...")
+    retries = 0
+    while retries < retry_times:
+        try:
+            request.urlretrieve(url, model_name, schedule)
+            break
+        except KeyboardInterrupt:
+            return False
+        except:
+            retries += 1
+            print(f"Download failed{', Retry downloading...' if retries < retry_times else '!'}")
+    return retries < retry_times
+
+
+def export_model(input_model, output_model, dataset_location):
+    print("\nexport model...")
+
+    if not os.path.exists("./table-transformer"):
+        subprocess.run("bash prepare.sh", shell=True)
+    
+    model_load_path = os.path.abspath(MODEL_URLS[input_model].split("/")[-1])
+    output_model = os.path.join(os.path.dirname(model_load_path), output_model)
+    if input_model == "detection_detr":
+        data_root_dir = os.path.join(dataset_location, "PubTables-1M-Detection")
+        data_type = "detection"
+        config_file = "detection_config.json"
+    elif input_model == "structure_detr":
+        data_root_dir = os.path.join(dataset_location, "PubTables-1M-Structure")
+        data_type = "structure"
+        config_file = "structure_config.json"
+    table_words_dir = os.path.join(data_root_dir, "words")
+
+    os.chdir("table-transformer/src")
+
+    command = f"python main.py \
+                --model_load_path {model_load_path} \
+                --output_model {output_model} \
+                --data_root_dir {data_root_dir} \
+                --table_words_dir {table_words_dir} \
+                --mode export \
+                --data_type {data_type} \
+                --device cpu \
+                --config_file {config_file}"
+
+    subprocess.run(command, shell=True)
+    assert os.path.exists(output_model), f"Export failed! {output_model} doesn't exist!"
+
+
+def prepare_model(input_model, output_model, dataset_location):
+    is_download_successful = download_model(MODEL_URLS[args.input_model], MAX_TIMES_RETRY_DOWNLOAD)
+    if is_download_successful:
+        export_model(input_model, output_model, dataset_location)
+
+
+if __name__ == "__main__":
+    args = parse_arguments()
+    prepare_model(args.input_model, args.output_model, args.dataset_location)

examples/onnxrt/object_detection/table_transformer/quantization/ptq_static/prepare_model.py

@@ -35,15 +35,28 @@ function run_benchmark {
 	      bash prepare.sh
     fi
 
+    if [[ "${input_model}" =~ "structure" ]]; then
+        task_data_dir="PubTables-1M-Structure"
+        data_type="structure"
+        config_file="structure_config.json"
+    fi
+    if [[ "${input_model}" =~ "detection" ]]; then
+        task_data_dir="PubTables-1M-Detection"
+        data_type="detection"
+        config_file="detection_config.json"
+    fi
+
+    input_model=$(realpath "$input_model")
+
     cd table-transformer/src
     python main.py \
             --input_onnx_model ${input_model} \
-            --data_root_dir ${dataset_location} \
-            --table_words_dir ${dataset_location}/words \
+            --data_root_dir "${dataset_location}/${task_data_dir}" \
+            --table_words_dir "${dataset_location}/${task_data_dir}/words" \
             --mode ${mode} \
-            --data_type structure \
+            --data_type ${data_type} \
             --device cpu \
-            --config_file structure_config.json
+            --config_file ${config_file}
 }
 
 main "$@"

examples/onnxrt/object_detection/table_transformer/quantization/ptq_static/run_benchmark.sh

@@ -35,16 +35,30 @@ function run_tuning {
 	      bash prepare.sh
     fi
 
+    if [[ "${input_model}" =~ "structure" ]]; then
+        task_data_dir="PubTables-1M-Structure"
+        data_type="structure"
+        config_file="structure_config.json"
+    fi
+    if [[ "${input_model}" =~ "detection" ]]; then
+        task_data_dir="PubTables-1M-Detection"
+        data_type="detection"
+        config_file="detection_config.json"
+    fi
+
+    input_model=$(realpath "$input_model")
+    output_model=$(realpath "$output_model")
+
     cd table-transformer/src
     python main.py \
             --input_onnx_model ${input_model} \
             --output_model ${output_model} \
-            --data_root_dir ${dataset_location} \
-            --table_words_dir ${dataset_location}/words \
+            --data_root_dir "${dataset_location}/${task_data_dir}" \
+            --table_words_dir "${dataset_location}/${task_data_dir}/words" \
             --mode quantize \
-            --data_type structure \
+            --data_type ${data_type} \
             --device cpu \
-            --config_file structure_config.json
+            --config_file ${config_file}
 }
 
 main "$@"