Merge pull request #966 from opendatalab/dev

Dev to 0.9.3

Dockerfile

@@ -42,7 +42,7 @@ RUN /bin/bash -c "wget https://gitee.com/myhloli/MinerU/raw/master/magic-pdf.tem
 
 # Download models and update the configuration file
 RUN /bin/bash -c "pip3 install modelscope && \
-    wget https://gitee.com/myhloli/MinerU/raw/master/docs/download_models.py && \
+    wget https://gitee.com/myhloli/MinerU/raw/master/scripts/download_models.py && \
     python3 download_models.py && \
     sed -i 's|/tmp/models|/root/.cache/modelscope/hub/opendatalab/PDF-Extract-Kit/models|g' /root/magic-pdf.json && \
     sed -i 's|cpu|cuda|g' /root/magic-pdf.json"

README.md

@@ -20,6 +20,7 @@
 [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/gist/myhloli/3b3a00a4a0a61577b6c30f989092d20d/mineru_demo.ipynb)
 [![Paper](https://img.shields.io/badge/Paper-arXiv-green)](https://arxiv.org/abs/2409.18839)
 
+
 <a href="https://trendshift.io/repositories/11174" target="_blank"><img src="https://trendshift.io/api/badge/repositories/11174" alt="opendatalab%2FMinerU | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
 
 <!-- language -->
@@ -41,6 +42,7 @@
 </div>
 
 # Changelog
+- 2024/11/15 0.9.3 released. Integrated [RapidTable](https://github.com/RapidAI/RapidTable) for table recognition, improving single-table parsing speed by more than 10 times, with higher accuracy and lower GPU memory usage.
 - 2024/11/06 0.9.2 released. Integrated the [StructTable-InternVL2-1B](https://huggingface.co/U4R/StructTable-InternVL2-1B) model for table recognition functionality.
 - 2024/10/31 0.9.0 released. This is a major new version with extensive code refactoring, addressing numerous issues, improving performance, reducing hardware requirements, and enhancing usability:
   - Refactored the sorting module code to use [layoutreader](https://github.com/ppaanngggg/layoutreader) for reading order sorting, ensuring high accuracy in various layouts.
@@ -75,10 +77,12 @@
             <ul>
             <li><a href="#online-demo">Online Demo</a></li>
             <li><a href="#quick-cpu-demo">Quick CPU Demo</a></li>
+            <li><a href="#using-gpu">Using GPU</a></li>
             </ul>
         </li>
         <li><a href="#usage">Usage</a>
             <ul>
+            <li><a href="#command-line">Command Line</a></li>
             <li><a href="#api">API</a></li>
             <li><a href="#deploy-derived-projects">Deploy Derived Projects</a></li>
             <li><a href="#development-guide">Development Guide</a></li>
@@ -87,6 +91,8 @@
       </ul>
     </li>
     <li><a href="#todo">TODO</a></li>
+    <li><a href="#known-issues">Known Issues</a></li>
+    <li><a href="#faq">FAQ</a></li>
     <li><a href="#all-thanks-to-our-contributors">All Thanks To Our Contributors</a></li>
     <li><a href="#license-information">License Information</a></li>
     <li><a href="#acknowledgments">Acknowledgments</a></li>
@@ -108,12 +114,89 @@ Compared to well-known commercial products, MinerU is still young. If you encoun
 
 https://github.com/user-attachments/assets/4bea02c9-6d54-4cd6-97ed-dff14340982c
 
+## Key Features
+
+- Remove headers, footers, footnotes, page numbers, etc., to ensure semantic coherence.
+- Output text in human-readable order, suitable for single-column, multi-column, and complex layouts.
+- Preserve the structure of the original document, including headings, paragraphs, lists, etc.
+- Extract images, image descriptions, tables, table titles, and footnotes.
+- Automatically recognize and convert formulas in the document to LaTeX format.
+- Automatically recognize and convert tables in the document to LaTeX or HTML format.
+- Automatically detect scanned PDFs and garbled PDFs and enable OCR functionality.
+- OCR supports detection and recognition of 84 languages.
+- Supports multiple output formats, such as multimodal and NLP Markdown, JSON sorted by reading order, and rich intermediate formats.
+- Supports various visualization results, including layout visualization and span visualization, for efficient confirmation of output quality.
+- Supports both CPU and GPU environments.
+- Compatible with Windows, Linux, and Mac platforms.
+
 ## Quick Start
 
-There are multiple different ways to experience MinerU:
+If you encounter any installation issues, please first consult the <a href="#faq">FAQ</a>. </br>
+If the parsing results are not as expected, refer to the <a href="#known-issues">Known Issues</a>. </br>
+There are three different ways to experience MinerU:
 
 - [Online Demo (No Installation Required)](#online-demo)
 - [Quick CPU Demo (Windows, Linux, Mac)](#quick-cpu-demo)
+- [Linux/Windows + CUDA](#Using-GPU)
+
+> [!WARNING]
+> **Pre-installation Notice—Hardware and Software Environment Support**
+>
+> To ensure the stability and reliability of the project, we only optimize and test for specific hardware and software environments during development. This ensures that users deploying and running the project on recommended system configurations will get the best performance with the fewest compatibility issues.
+>
+> By focusing resources on the mainline environment, our team can more efficiently resolve potential bugs and develop new features.
+>
+> In non-mainline environments, due to the diversity of hardware and software configurations, as well as third-party dependency compatibility issues, we cannot guarantee 100% project availability. Therefore, for users who wish to use this project in non-recommended environments, we suggest carefully reading the documentation and FAQ first. Most issues already have corresponding solutions in the FAQ. We also encourage community feedback to help us gradually expand support.
+
+<table>
+    <tr>
+        <td colspan="3" rowspan="2">Operating System</td>
+    </tr>
+    <tr>
+        <td>Ubuntu 22.04 LTS</td>
+        <td>Windows 10 / 11</td>
+        <td>macOS 11+</td>
+    </tr>
+    <tr>
+        <td colspan="3">CPU</td>
+        <td>x86_64(unsupported ARM Linux)</td>
+        <td>x86_64(unsupported ARM Windows)</td>
+        <td>x86_64 / arm64</td>
+    </tr>
+    <tr>
+        <td colspan="3">Memory</td>
+        <td colspan="3">16GB or more, recommended 32GB+</td>
+    </tr>
+    <tr>
+        <td colspan="3">Python Version</td>
+        <td colspan="3">3.10(Please make sure to create a Python 3.10 virtual environment using conda)</td>
+    </tr>
+    <tr>
+        <td colspan="3">Nvidia Driver Version</td>
+        <td>latest (Proprietary Driver)</td>
+        <td>latest</td>
+        <td>None</td>
+    </tr>
+    <tr>
+        <td colspan="3">CUDA Environment</td>
+        <td>Automatic installation [12.1 (pytorch) + 11.8 (paddle)]</td>
+        <td>11.8 (manual installation) + cuDNN v8.7.0 (manual installation)</td>
+        <td>None</td>
+    </tr>
+    <tr>
+        <td rowspan="2">GPU Hardware Support List</td>
+        <td colspan="2">Minimum Requirement 8G+ VRAM</td>
+        <td colspan="2">3060ti/3070/4060<br>
+        8G VRAM enables layout, formula recognition acceleration and OCR acceleration</td>
+        <td rowspan="2">None</td>
+    </tr>
+    <tr>
+        <td colspan="2">Recommended Configuration 10G+ VRAM</td>
+        <td colspan="2">3080/3080ti/3090/3090ti/4070/4070ti/4070tisuper/4080/4090<br>
+        10G VRAM or more can enable layout, formula recognition, OCR acceleration and table recognition acceleration simultaneously
+        </td>
+    </tr>
+</table>
 
 ### Online Demo
 
@@ -148,6 +231,7 @@ You can find the `magic-pdf.json` file in your 【user directory】.
 
 You can modify certain configurations in this file to enable or disable features, such as table recognition:
 
+
 > [!NOTE]
 > If the following items are not present in the JSON, please manually add the required items and remove the comment content (standard JSON does not support comments).
 
@@ -163,15 +247,92 @@ You can modify certain configurations in this file to enable or disable features
         "enable": true  // The formula recognition feature is enabled by default. If you need to disable it, please change the value here to "false".
     },
     "table-config": {
-        "model": "tablemaster",  // When using structEqTable, please change to "struct_eqtable".
+        "model": "rapid_table",  // When using structEqTable, please change to "struct_eqtable".
         "enable": false, // The table recognition feature is disabled by default. If you need to enable it, please change the value here to "true".
         "max_time": 400
     }
 }
 ```
 
+### Using GPU
+
+If your device supports CUDA and meets the GPU requirements of the mainline environment, you can use GPU acceleration. Please select the appropriate guide based on your system:
+
+- [Ubuntu 22.04 LTS + GPU](docs/README_Ubuntu_CUDA_Acceleration_en_US.md)
+- [Windows 10/11 + GPU](docs/README_Windows_CUDA_Acceleration_en_US.md)
+- Quick Deployment with Docker
+> [!IMPORTANT]
+> Docker requires a GPU with at least 8GB of VRAM, and all acceleration features are enabled by default.
+>
+> Before running this Docker, you can use the following command to check if your device supports CUDA acceleration on Docker.
+> 
+> ```bash
+> docker run --rm --gpus=all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi
+> ```
+  ```bash
+  wget https://github.com/opendatalab/MinerU/raw/master/Dockerfile
+  docker build -t mineru:latest .
+  docker run --rm -it --gpus=all mineru:latest /bin/bash
+  magic-pdf --help
+  ```
+
 ## Usage
 
+### Command Line
+
+```bash
+magic-pdf --help
+Usage: magic-pdf [OPTIONS]
+
+Options:
+  -v, --version                display the version and exit
+  -p, --path PATH              local pdf filepath or directory  [required]
+  -o, --output-dir PATH        output local directory  [required]
+  -m, --method [ocr|txt|auto]  the method for parsing pdf. ocr: using ocr
+                               technique to extract information from pdf. txt:
+                               suitable for the text-based pdf only and
+                               outperform ocr. auto: automatically choose the
+                               best method for parsing pdf from ocr and txt.
+                               without method specified, auto will be used by
+                               default.
+  -l, --lang TEXT              Input the languages in the pdf (if known) to
+                               improve OCR accuracy.  Optional. You should
+                               input "Abbreviation" with language form url: ht
+                               tps://paddlepaddle.github.io/PaddleOCR/latest/en
+                               /ppocr/blog/multi_languages.html#5-support-languages-
+                               and-abbreviations
+  -d, --debug BOOLEAN          Enables detailed debugging information during
+                               the execution of the CLI commands.
+  -s, --start INTEGER          The starting page for PDF parsing, beginning
+                               from 0.
+  -e, --end INTEGER            The ending page for PDF parsing, beginning from
+                               0.
+  --help                       Show this message and exit.
+
+
+## show version
+magic-pdf -v
+
+## command line example
+magic-pdf -p {some_pdf} -o {some_output_dir} -m auto
+```
+
+`{some_pdf}` can be a single PDF file or a directory containing multiple PDFs.
+The results will be saved in the `{some_output_dir}` directory. The output file list is as follows:
+
+```text
+├── some_pdf.md                          # markdown file
+├── images                               # directory for storing images
+├── some_pdf_layout.pdf                  # layout diagram (Include layout reading order)
+├── some_pdf_middle.json                 # MinerU intermediate processing result
+├── some_pdf_model.json                  # model inference result
+├── some_pdf_origin.pdf                  # original PDF file
+├── some_pdf_spans.pdf                   # smallest granularity bbox position information diagram
+└── some_pdf_content_list.json           # Rich text JSON arranged in reading order
+```
+> [!TIP]
+> For more information about the output files, please refer to the [Output File Description](docs/output_file_en_us.md).
+
 ### API
 
 Processing files from local disk
@@ -228,6 +389,24 @@ TODO
 - [ ] [Chemical formula recognition](docs/chemical_knowledge_introduction/introduction.pdf)
 - [ ] Geometric shape recognition
 
+# Known Issues
+
+- Reading order is determined by the model based on the spatial distribution of readable content, and may be out of order in some areas under extremely complex layouts.
+- Vertical text is not supported.
+- Tables of contents and lists are recognized through rules, and some uncommon list formats may not be recognized.
+- Only one level of headings is supported; hierarchical headings are not currently supported.
+- Code blocks are not yet supported in the layout model.
+- Comic books, art albums, primary school textbooks, and exercises cannot be parsed well.
+- Table recognition may result in row/column recognition errors in complex tables.
+- OCR recognition may produce inaccurate characters in PDFs of lesser-known languages (e.g., diacritical marks in Latin script, easily confused characters in Arabic script).
+- Some formulas may not render correctly in Markdown.
+
+# FAQ
+
+[FAQ in Chinese](docs/FAQ_zh_cn.md)
+
+[FAQ in English](docs/FAQ_en_us.md)
+
 # All Thanks To Our Contributors
 
 <a href="https://github.com/opendatalab/MinerU/graphs/contributors">

README_ja-JP.md

+> [!Warning]
+> このドキュメントはすでに古くなっています。最新版のドキュメントを参照してください：[ENGLISH](README.md)。
 <div id="top">
 
 <p align="center">
@@ -18,9 +20,7 @@
 <a href="https://trendshift.io/repositories/11174" target="_blank"><img src="https://trendshift.io/api/badge/repositories/11174" alt="opendatalab%2FMinerU | Trendshift" style="width: 200px; height: 55px;"/></a>
 
 
-<div align="center" style="color: red; background-color: #ffdddd; padding: 10px; border: 1px solid red; border-radius: 5px;">
-  <strong>NOTE：</strong> このドキュメントはすでに古くなっています。最新版のドキュメントを参照してください。
-</div>
+
 
 
 [English](README.md) | [简体中文](README_zh-CN.md) | [日本語](README_ja-JP.md)

demo/magic_pdf_parse_main.py

@@ -19,9 +19,10 @@ def json_md_dump(
         pdf_name,
         content_list,
         md_content,
+        orig_model_list,
 ):
     # 写入模型结果到 model.json
-    orig_model_list = copy.deepcopy(pipe.model_list)
+
     md_writer.write(
         content=json.dumps(orig_model_list, ensure_ascii=False, indent=4),
         path=f"{pdf_name}_model.json"
@@ -87,9 +88,12 @@ def pdf_parse_main(
 
         pdf_bytes = open(pdf_path, "rb").read()  # 读取 pdf 文件的二进制数据
 
+        orig_model_list = []
+
         if model_json_path:
             # 读取已经被模型解析后的pdf文件的 json 原始数据，list 类型
             model_json = json.loads(open(model_json_path, "r", encoding="utf-8").read())
+            orig_model_list = copy.deepcopy(model_json)
         else:
             model_json = []
 
@@ -115,8 +119,9 @@ def pdf_parse_main(
         pipe.pipe_classify()
 
         # 如果没有传入模型数据，则使用内置模型解析
-        if not model_json:
+        if len(model_json) == 0:
             pipe.pipe_analyze()  # 解析
+            orig_model_list = copy.deepcopy(pipe.model_list)
 
         # 执行解析
         pipe.pipe_parse()
@@ -126,7 +131,7 @@ def pdf_parse_main(
         md_content = pipe.pipe_mk_markdown(image_path_parent, drop_mode="none")
 
         if is_json_md_dump:
-            json_md_dump(pipe, md_writer, pdf_name, content_list, md_content)
+            json_md_dump(pipe, md_writer, pdf_name, content_list, md_content, orig_model_list)
 
         if is_draw_visualization_bbox:
             draw_visualization_bbox(pipe.pdf_mid_data['pdf_info'], pdf_bytes, output_path, pdf_name)

docs/download_models.py

-import json
-import os
-
-import requests
-from modelscope import snapshot_download
-
-
-def download_json(url):
-    # 下载JSON文件
-    response = requests.get(url)
-    response.raise_for_status()  # 检查请求是否成功
-    return response.json()
-
-
-def download_and_modify_json(url, local_filename, modifications):
-    if os.path.exists(local_filename):
-        data = json.load(open(local_filename))
-        config_version = data.get('config_version', '0.0.0')
-        if config_version < '1.0.0':
-            data = download_json(url)
-    else:
-        data = download_json(url)
-
-    # 修改内容
-    for key, value in modifications.items():
-        data[key] = value
-
-    # 保存修改后的内容
-    with open(local_filename, 'w', encoding='utf-8') as f:
-        json.dump(data, f, ensure_ascii=False, indent=4)
-
-
-if __name__ == '__main__':
-    mineru_patterns = [
-        "models/Layout/LayoutLMv3/*",
-        "models/Layout/YOLO/*",
-        "models/MFD/YOLO/*",
-        "models/MFR/unimernet_small/*",
-        "models/TabRec/TableMaster/*",
-        "models/TabRec/StructEqTable/*",
-    ]
-    model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', allow_patterns=mineru_patterns)
-    layoutreader_model_dir = snapshot_download('ppaanngggg/layoutreader')
-    model_dir = model_dir + '/models'
-    print(f'model_dir is: {model_dir}')
-    print(f'layoutreader_model_dir is: {layoutreader_model_dir}')
-
-    json_url = 'https://gitee.com/myhloli/MinerU/raw/dev/magic-pdf.template.json'
-    config_file_name = 'magic-pdf.json'
-    home_dir = os.path.expanduser('~')
-    config_file = os.path.join(home_dir, config_file_name)
-
-    json_mods = {
-        'models-dir': model_dir,
-        'layoutreader-model-dir': layoutreader_model_dir,
-    }
-
-    download_and_modify_json(json_url, config_file, json_mods)
-    print(f'The configuration file has been configured successfully, the path is: {config_file}')

docs/download_models_hf.py

-import json
-import os
-
-import requests
-from huggingface_hub import snapshot_download
-
-
-def download_json(url):
-    # 下载JSON文件
-    response = requests.get(url)
-    response.raise_for_status()  # 检查请求是否成功
-    return response.json()
-
-
-def download_and_modify_json(url, local_filename, modifications):
-    if os.path.exists(local_filename):
-        data = json.load(open(local_filename))
-        config_version = data.get('config_version', '0.0.0')
-        if config_version < '1.0.0':
-            data = download_json(url)
-    else:
-        data = download_json(url)
-
-    # 修改内容
-    for key, value in modifications.items():
-        data[key] = value
-
-    # 保存修改后的内容
-    with open(local_filename, 'w', encoding='utf-8') as f:
-        json.dump(data, f, ensure_ascii=False, indent=4)
-
-
-if __name__ == '__main__':
-
-    mineru_patterns = [
-        "models/Layout/LayoutLMv3/*",
-        "models/Layout/YOLO/*",
-        "models/MFD/YOLO/*",
-        "models/MFR/unimernet_small/*",
-        "models/TabRec/TableMaster/*",
-        "models/TabRec/StructEqTable/*",
-    ]
-    model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', allow_patterns=mineru_patterns)
-
-    layoutreader_pattern = [
-        "*.json",
-        "*.safetensors",
-    ]
-    layoutreader_model_dir = snapshot_download('hantian/layoutreader', allow_patterns=layoutreader_pattern)
-
-    model_dir = model_dir + '/models'
-    print(f'model_dir is: {model_dir}')
-    print(f'layoutreader_model_dir is: {layoutreader_model_dir}')
-
-    json_url = 'https://github.com/opendatalab/MinerU/raw/dev/magic-pdf.template.json'
-    config_file_name = 'magic-pdf.json'
-    home_dir = os.path.expanduser('~')
-    config_file = os.path.join(home_dir, config_file_name)
-
-    json_mods = {
-        'models-dir': model_dir,
-        'layoutreader-model-dir': layoutreader_model_dir,
-    }
-
-    download_and_modify_json(json_url, config_file, json_mods)
-    print(f'The configuration file has been configured successfully, the path is: {config_file}')

docs/how_to_download_models_en.md

@@ -8,7 +8,7 @@ Use a Python Script to Download Model Files from Hugging Face
 
 ```bash
 pip install huggingface_hub
-wget https://github.com/opendatalab/MinerU/raw/master/docs/download_models_hf.py -O download_models_hf.py
+wget https://github.com/opendatalab/MinerU/raw/master/scripts/download_models_hf.py -O download_models_hf.py
 python download_models_hf.py
 ```
 

docs/how_to_download_models_zh_cn.md

@@ -8,7 +8,7 @@
   <summary>方法一：从 Hugging Face 下载模型</summary>
   <p>使用python脚本 从Hugging Face下载模型文件</p>
   <pre><code>pip install huggingface_hub
-wget https://gitee.com/myhloli/MinerU/raw/master/docs/download_models_hf.py -O download_models_hf.py
+wget https://gitee.com/myhloli/MinerU/raw/master/scripts/download_models_hf.py -O download_models_hf.py
 python download_models_hf.py</code></pre>
 </details>
 
@@ -18,7 +18,7 @@ python download_models_hf.py</code></pre>
 
 ```bash
 pip install modelscope
-wget https://gitee.com/myhloli/MinerU/raw/master/docs/download_models.py -O download_models.py
+wget https://gitee.com/myhloli/MinerU/raw/master/scripts/download_models.py -O download_models.py
 python download_models.py
 ```
 

magic-pdf.template.json

@@ -15,7 +15,7 @@
         "enable": true
     },
     "table-config": {
-        "model": "tablemaster",
+        "model": "rapid_table",
         "enable": false,
         "max_time": 400
     },

magic_pdf/dict2md/ocr_mkcontent.py

@@ -168,7 +168,7 @@ def merge_para_with_text(para_block):
                         # 如果是前一行带有-连字符，那么末尾不应该加空格
                         if __is_hyphen_at_line_end(content):
                             para_text += content[:-1]
-                        elif len(content) == 1 and content not in ['A', 'I', 'a', 'i']:
+                        elif len(content) == 1 and content not in ['A', 'I', 'a', 'i'] and not content.isdigit():
                             para_text += content
                         else:  # 西方文本语境下 content间需要空格分隔
                             para_text += f"{content} "

magic_pdf/libs/Constants.py

@@ -50,4 +50,6 @@ class MODEL_NAME:
 
     YOLO_V8_MFD = "yolo_v8_mfd"
 
-    UniMerNet_v2_Small = "unimernet_small"
+    UniMerNet_v2_Small = "unimernet_small"
\ No newline at end of file
+
+    RAPID_TABLE = "rapid_table"
\ No newline at end of file

magic_pdf/libs/config_reader.py

@@ -92,7 +92,7 @@ def get_table_recog_config():
     table_config = config.get('table-config')
     if table_config is None:
         logger.warning(f"'table-config' not found in {CONFIG_FILE_NAME}, use 'False' as default")
-        return json.loads(f'{{"model": "{MODEL_NAME.TABLE_MASTER}","enable": false, "max_time": 400}}')
+        return json.loads(f'{{"model": "{MODEL_NAME.RAPID_TABLE}","enable": false, "max_time": 400}}')
     else:
         return table_config
 

magic_pdf/model/pek_sub_modules/post_process.py

-import re
-
-def layout_rm_equation(layout_res):
-    rm_idxs = []
-    for idx, ele in enumerate(layout_res['layout_dets']):
-        if ele['category_id'] == 10:
-            rm_idxs.append(idx)
-    
-    for idx in rm_idxs[::-1]:
-        del layout_res['layout_dets'][idx]
-    return layout_res
-
-
-def get_croped_image(image_pil, bbox):
-    x_min, y_min, x_max, y_max = bbox
-    croped_img = image_pil.crop((x_min, y_min, x_max, y_max))
-    return croped_img
-
-
-def latex_rm_whitespace(s: str):
-    """Remove unnecessary whitespace from LaTeX code.
-    """
-    text_reg = r'(\\(operatorname|mathrm|text|mathbf)\s?\*? {.*?})'
-    letter = '[a-zA-Z]'
-    noletter = '[\W_^\d]'
-    names = [x[0].replace(' ', '') for x in re.findall(text_reg, s)]
-    s = re.sub(text_reg, lambda match: str(names.pop(0)), s)
-    news = s
-    while True:
-        s = news
-        news = re.sub(r'(?!\\ )(%s)\s+?(%s)' % (noletter, noletter), r'\1\2', s)
-        news = re.sub(r'(?!\\ )(%s)\s+?(%s)' % (noletter, letter), r'\1\2', news)
-        news = re.sub(r'(%s)\s+?(%s)' % (letter, noletter), r'\1\2', news)
-        if news == s:
-            break
-    return s
\ No newline at end of file

magic_pdf/model/sub_modules/layout/doclayout_yolo/DocLayoutYOLO.py

+from doclayout_yolo import YOLOv10
+
+
+class DocLayoutYOLOModel(object):
+    def __init__(self, weight, device):
+        self.model = YOLOv10(weight)
+        self.device = device
+
+    def predict(self, image):
+        layout_res = []
+        doclayout_yolo_res = self.model.predict(image, imgsz=1024, conf=0.25, iou=0.45, verbose=True, device=self.device)[0]
+        for xyxy, conf, cla in zip(doclayout_yolo_res.boxes.xyxy.cpu(), doclayout_yolo_res.boxes.conf.cpu(),
+                                   doclayout_yolo_res.boxes.cls.cpu()):
+            xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
+            new_item = {
+                'category_id': int(cla.item()),
+                'poly': [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
+                'score': round(float(conf.item()), 3),
+            }
+            layout_res.append(new_item)
+        return layout_res
\ No newline at end of file