Merge pull request #2 from UniModal4Reasoning/dev

update README.md and demo.py

Author: PrinceVictor

README.md

@@ -15,12 +15,15 @@ Welcome to the official repository of StructEqTable-Deploy, a solution that conv
 ## Abstract
 Table is an effective way to represent structured data in scientific publications, financial statements, invoices, web pages, and many other scenarios. Extracting tabular data from a visual table image and performing the downstream reasoning tasks according to the extracted data is challenging, mainly due to that tables often present complicated column and row headers with spanning cell operation. To address these challenges, we present TableX, a large-scale multi-modal table benchmark extracted from [DocGenome benchmark](https://unimodal4reasoning.github.io/DocGenome_page/) for table pre-training, comprising more than 2 million high-quality Image-LaTeX pair data covering 156 disciplinary classes. Besides, benefiting from such large-scale data, we train an end-to-end model, StructEqTable, which provides the capability to precisely obtain the corresponding LaTeX description from a visual table image and perform multiple table-related reasoning tasks, including structural extraction and question answering, broadening its application scope and potential.
 
+## Release
+- [2024/7/30] 🔥 We have released the first version of StructEqTable. (Current version of StructEqTable is able to process table images from scientific documents such as arXiv, Scihub papers. Times New Roman And Songti(宋体) are main fonts used in table image, other fonts may decrease the accuracy of the model's output.)
 
 ## TODO
 
 - [x] Release inference code and checkpoints of StructEqTable.
 - [x] Support Chinese version of StructEqTable.
-- [ ] Improve the inference speed of StructEqTable.
+- [ ] Expand more domains of table image to improve the model's general capabilities.
+- [ ] Accelerated version of StructEqTable using TensorRT-LLM.
 
 
 ### Installation
@@ -35,14 +38,35 @@ pip install "git+https://github.com/UniModal4Reasoning/StructEqTable-Deploy.git"
 ```
 
 ## Quick Demo
-- run the demo/demo.py
+- Run the demo/demo.py
 ```shell script
 cd demo
 
-python demo.py \ --image_path ./demo.png \
-  --ckpt_path ${CKPT_PATH}
+python demo.py \
+  --image_path ./demo.png \
+  --ckpt_path ${CKPT_PATH} \
+  --output_format latex
 ```
 
+- Obtain other format output
+
+  Our model output Latex format code by default. 
+  If you want to get other format like HTML or Markdown, 
+  `pypandoc` support convert latex format code into HTML and Markdown format for simple table (table has no merge cell ).
+
+```shell script
+sudo apt install pandoc
+pip install pypandoc
+
+cd demo
+
+python demo.py \
+  --image_path ./demo.png \
+  --ckpt_path ${CKPT_PATH} \
+  --output_format html markdown
+```
+
+
 - Visualization Results
     - The input data are sampled from SciHub domain.
 

demo/demo.py

@@ -10,18 +10,20 @@ def parse_config():
     parser = argparse.ArgumentParser(description='arg parser')
     parser.add_argument('--image_path', type=str, default='demo.png', help='data path for table image')
     parser.add_argument('--ckpt_path', type=str, default='U4R/StructTable-base', help='ckpt path for table model, which can be downloaded from huggingface')
+    parser.add_argument('-t', '--max_waiting_time', type=int, default=60, help='maximum waiting time of model inference')
     parser.add_argument('--cpu', action='store_true', default=False, help='using cpu for inference')
-    parser.add_argument('--html', action='store_true', default=False, help='output html format table code')
+    parser.add_argument('-f', '--output_format', type=str, nargs='+', default=['latex'], 
+                        help='The model outputs LaTeX format code by default. Simple structured table LaTeX code can be converted to HTML or Markdown format using pypandoc.')
     args = parser.parse_args()
     return args
 
 def main():
     args = parse_config()
-    if args.html:
+    if 'html' in args.output_format or 'markdown' in args.output_format:
         from pypandoc import convert_text
 
     # build model
-    model = build_model(args.ckpt_path, max_new_tokens=4096, max_time=60)
+    model = build_model(args.ckpt_path, max_new_tokens=4096, max_time=args.max_waiting_time)
     if not args.cpu:
         model = model.cuda()
 
@@ -35,12 +37,17 @@ def main():
     # show output latex code of table
     cost_time = time.time() - start_time
     print(f"total cost time: {cost_time:.2f}s")
+
+    if cost_time >= args.max_waiting_time:
+        warn_log = f"\033[93mThe model inference time exceeds the maximum waiting time {args.max_waiting_time} seconds, the result may be incomplete.\n" \
+        "Please increase the maximum waiting time with argument --max_waiting_time or Model may not support the type of input table image \033[0m"
+        print(warn_log)
+
+    
     for i, latex_code in enumerate(output):
-        if args.html:
-            html_code = convert_text(latex_code, 'html', format='latex')
-            print(f"Table {i} HTML code:\n{html_code}")
-        else:
-            print(f"Table {i} LaTex code:\n{latex_code}")
+        for tgt_fmt in args.output_format:
+            tgt_code = convert_text(latex_code, tgt_fmt, format='latex') if tgt_fmt != 'latex' else latex_code
+            print(f"Table {i} {tgt_fmt.upper()} format output:\n{tgt_code}")
 
 
 if __name__ == '__main__':