Skip to content

GitLab

Unverified Commit 1f497129 authored by Xiaomeng Zhao's avatar Xiaomeng Zhao Committed by GitHub
Browse files

Merge pull request #1759 from opendatalab/release-1.2.0 

Release 1.2.0
parents 9a87d3ea bd3a7b37
Hide whitespace changes
Inline Side-by-side
Showing with 513 additions and 298 deletions
magic_pdf/pdf_parse_union_core_v2.py
View file @ 1f497129
......@@ -6,8 +6,10 @@ import statistics
import time
from typing import List
import cv2
import fitz
import torch
import numpy as np
from loguru import logger
from magic_pdf.config.enums import SupportedPdfParseMethod
......@@ -127,16 +129,15 @@ def fill_char_in_spans(spans, all_chars):
span['chars'].append(char)
break
empty_spans = []
need_ocr_spans = []
for span in spans:
chars_to_content(span)
# 有的span中虽然没有字但有一两个空的占位符，用宽高和content长度过滤
if len(span['content']) * span['height'] < span['width'] * 0.5:
# logger.info(f"maybe empty span: {len(span['content'])}, {span['height']}, {span['width']}")
empty_spans.append(span)
need_ocr_spans.append(span)
del span['height'], span['width']
return empty_spans
return need_ocr_spans
# 使用鲁棒性更强的中心点坐标判断
......@@ -190,6 +191,31 @@ def remove_tilted_line(text_blocks):
block['lines'].remove(line)
def calculate_contrast(img, img_mode) -> float:
"""
计算给定图像的对比度。
:param img: 图像，类型为numpy.ndarray
:Param img_mode = 图像的色彩通道，'rgb' 或 'bgr'
:return: 图像的对比度值
"""
if img_mode == 'rgb':
# 将RGB图像转换为灰度图
gray_img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
elif img_mode == 'bgr':
# 将BGR图像转换为灰度图
gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
else:
raise ValueError("Invalid image mode. Please provide 'rgb' or 'bgr'.")
# 计算均值和标准差
mean_value = np.mean(gray_img)
std_dev = np.std(gray_img)
# 对比度定义为标准差除以平均值（加上小常数避免除零错误）
contrast = std_dev / (mean_value + 1e-6)
# logger.info(f"contrast: {contrast}")
return round(contrast, 2)
def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang):
# cid用0xfffd表示，连字符拆开
# text_blocks_raw = pdf_page.get_text('rawdict', flags=fitz.TEXT_PRESERVE_WHITESPACE | fitz.TEXT_MEDIABOX_CLIP)['blocks']
......@@ -274,9 +300,9 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang
span['chars'] = []
new_spans.append(span)
empty_spans = fill_char_in_spans(new_spans, all_pymu_chars)
need_ocr_spans = fill_char_in_spans(new_spans, all_pymu_chars)
if len(empty_spans) > 0:
if len(need_ocr_spans) > 0:
# 初始化ocr模型
atom_model_manager = AtomModelSingleton()
......@@ -287,9 +313,15 @@ def txt_spans_extract_v2(pdf_page, spans, all_bboxes, all_discarded_blocks, lang
lang=lang
)
for span in empty_spans:
for span in need_ocr_spans:
# 对span的bbox截图再ocr
span_img = cut_image_to_pil_image(span['bbox'], pdf_page, mode='cv2')
# 计算span的对比度，低于0.20的span不进行ocr
if calculate_contrast(span_img, img_mode='bgr') <= 0.20:
spans.remove(span)
continue
ocr_res = ocr_model.ocr(span_img, det=False)
if ocr_res and len(ocr_res) > 0:
if len(ocr_res[0]) > 0:
......@@ -404,10 +436,11 @@ def cal_block_index(fix_blocks, sorted_bboxes):
block_bboxes.append(block['bbox'])
# 删除图表body block中的虚拟line信息, 并用real_lines信息回填
if block['type'] in [BlockType.ImageBody, BlockType.TableBody]:
block['virtual_lines'] = copy.deepcopy(block['lines'])
block['lines'] = copy.deepcopy(block['real_lines'])
del block['real_lines']
if block['type'] in [BlockType.ImageBody, BlockType.TableBody, BlockType.Title, BlockType.InterlineEquation]:
if 'real_lines' in block:
block['virtual_lines'] = copy.deepcopy(block['lines'])
block['lines'] = copy.deepcopy(block['real_lines'])
del block['real_lines']
import numpy as np
......
magic_pdf/post_proc/llm_aided.py
View file @ 1f497129
......@@ -3,6 +3,7 @@ import json
from loguru import logger
from magic_pdf.dict2md.ocr_mkcontent import merge_para_with_text
from openai import OpenAI
import ast
#@todo: 有的公式以"\"结尾，这样会导致尾部拼接的"$"被转义，也需要修复
......@@ -119,11 +120,12 @@ def llm_aided_title(pdf_info_dict, title_aided_config):
- 在完成初步分级后，仔细检查分级结果的合理性
- 根据上下文关系和逻辑顺序，对不合理的分级进行微调
- 确保最终的分级结果符合文档的实际结构和逻辑
- 字典中可能包含被误当成标题的正文，你可以通过将其层级标记为 0 来排除它们
IMPORTANT:
请直接返回优化过的由标题层级组成的json，格式如下：
{{"0":1,"1":2,"2":2,"3":3}}
返回的json不需要格式化。
请直接返回优化过的由标题层级组成的字典，格式为{{标题id:标题层级}}，如下：
{{0:1,1:2,2:2,3:3}}
不需要对字典格式化，不需要返回任何其他信息。
Input title list:
{title_dict}
......@@ -133,7 +135,7 @@ Corrected title list:
retry_count = 0
max_retries = 3
json_completion = None
dict_completion = None
while retry_count < max_retries:
try:
......@@ -143,24 +145,20 @@ Corrected title list:
{'role': 'user', 'content': title_optimize_prompt}],
temperature=0.7,
)
json_completion = json.loads(completion.choices[0].message.content)
# logger.info(f"Title completion: {completion.choices[0].message.content}")
dict_completion = ast.literal_eval(completion.choices[0].message.content)
# logger.info(f"len(dict_completion): {len(dict_completion)}, len(title_dict): {len(title_dict)}")
# logger.info(f"Title completion: {json_completion}")
# logger.info(f"len(json_completion): {len(json_completion)}, len(title_dict): {len(title_dict)}")
if len(json_completion) == len(title_dict):
if len(dict_completion) == len(title_dict):
for i, origin_title_block in enumerate(origin_title_list):
origin_title_block["level"] = int(json_completion[str(i)])
origin_title_block["level"] = int(dict_completion[i])
break
else:
logger.warning("The number of titles in the optimized result is not equal to the number of titles in the input.")
retry_count += 1
except Exception as e:
if isinstance(e, json.decoder.JSONDecodeError):
logger.warning(f"JSON decode error on attempt {retry_count + 1}: {e}")
else:
logger.exception(e)
logger.exception(e)
retry_count += 1
if json_completion is None:
logger.error("Failed to decode JSON after maximum retries.")
if dict_completion is None:
logger.error("Failed to decode dict after maximum retries.")
magic_pdf/pre_proc/ocr_dict_merge.py
View file @ 1f497129
......@@ -60,6 +60,19 @@ def merge_spans_to_line(spans, threshold=0.6):
return lines
def span_block_type_compatible(span_type, block_type):
if span_type in [ContentType.Text, ContentType.InlineEquation]:
return block_type in [BlockType.Text, BlockType.Title, BlockType.ImageCaption, BlockType.ImageFootnote, BlockType.TableCaption, BlockType.TableFootnote]
elif span_type == ContentType.InterlineEquation:
return block_type in [BlockType.InterlineEquation]
elif span_type == ContentType.Image:
return block_type in [BlockType.ImageBody]
elif span_type == ContentType.Table:
return block_type in [BlockType.TableBody]
else:
return False
def fill_spans_in_blocks(blocks, spans, radio):
"""将allspans中的span按位置关系，放入blocks中."""
block_with_spans = []
......@@ -78,8 +91,7 @@ def fill_spans_in_blocks(blocks, spans, radio):
block_spans = []
for span in spans:
span_bbox = span['bbox']
if calculate_overlap_area_in_bbox1_area_ratio(
span_bbox, block_bbox) > radio:
if calculate_overlap_area_in_bbox1_area_ratio(span_bbox, block_bbox) > radio and span_block_type_compatible(span['type'], block_type):
block_spans.append(span)
block_dict['spans'] = block_spans
......
projects/web_api/Dockerfile
View file @ 1f497129
# Use the official Ubuntu base image
FROM ubuntu:latest
FROM python:3.10-slim-bookworm AS base
# ENV http_proxy http://127.0.0.1:7890
# ENV https_proxy http://127.0.0.1:7890
WORKDIR /app
# Set environment variables to non-interactive to avoid prompts during installation
ENV DEBIAN_FRONTEND=noninteractive
ENV LANG C.UTF-8
# ADD sources.list /etc/apt
# RUN apt-get clean
ENV DEBIAN_FRONTEND=noninteractive \
LANG=C.UTF-8 \
PYTHONDONTWRITEBYTECODE=1 \
PYTHONUNBUFFERED=1 \
PIP_DISABLE_PIP_VERSION_CHECK=1 \
PIP_NO_CACHE_DIR=1
FROM base AS build
# Update the package list and install necessary packages
RUN apt-get -q update \
&& apt-get -q install -y --no-install-recommends \
apt-utils \
bats \
build-essential
RUN apt-get update && apt-get install -y vim net-tools procps lsof curl wget iputils-ping telnet lrzsz git
RUN apt-get update && \
apt-get install -y \
software-properties-common && \
add-apt-repository ppa:deadsnakes/ppa && \
apt-get update && \
apt-get install -y \
python3.10 \
python3.10-venv \
python3.10-distutils \
python3-pip \
wget \
git \
libgl1 \
libglib2.0-0 \
&& rm -rf /var/lib/apt/lists/*
# RUN unset http_proxy && unset https_proxy
# Set Python 3.10 as the default python3
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1
# Create a virtual environment for MinerU
RUN python3 -m venv /opt/mineru_venv
RUN pip config set global.index-url https://mirrors.aliyun.com/pypi/simple
# Activate the virtual environment and install necessary Python packages
RUN /bin/bash -c "source /opt/mineru_venv/bin/activate && \
pip install --upgrade pip && \
pip install magic-pdf[full] --extra-index-url https://myhloli.github.io/wheels/ --no-cache-dir"
apt-get install -y --no-install-recommends \
build-essential && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
# Build Python dependencies
COPY requirements.txt .
RUN python -m venv /app/venv && \
. /app/venv/bin/activate && \
pip install -r requirements.txt && \
pip uninstall -y paddlepaddle && \
pip install -i https://www.paddlepaddle.org.cn/packages/stable/cu118/ \
paddlepaddle-gpu==3.0.0rc1
RUN /bin/bash -c "source /opt/mineru_venv/bin/activate && \
pip install fastapi uvicorn python-multipart --no-cache-dir"
# Download models
COPY download_models.py .
RUN . /app/venv/bin/activate && \
./download_models.py
RUN /bin/bash -c "source /opt/mineru_venv/bin/activate && \
pip uninstall paddlepaddle -y"
RUN /bin/bash -c "source /opt/mineru_venv/bin/activate && \
python -m pip install paddlepaddle-gpu==3.0.0b1 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/ --no-cache-dir"
FROM base AS prod
# Copy the configuration file template and set up the model directory
COPY magic-pdf.template.json /root/magic-pdf.json
ADD models /opt/models
ADD .paddleocr /root/.paddleocr
ADD app.py /root/app.py
# Copy Python dependencies and models from the build stage
COPY --from=build /app/venv /app/venv
COPY --from=build /opt/models /opt/models
COPY --from=build /opt/layoutreader /opt/layoutreader
WORKDIR /root
# Set the models directory in the configuration file (adjust the path as needed)
RUN sed -i 's|/tmp/models|/opt/models|g' /root/magic-pdf.json
# Create the models directory
# RUN mkdir -p /opt/models
# Update the package list and install necessary packages
RUN apt-get update && \
apt-get install -y --no-install-recommends \
libgl1 \
libglib2.0-0 \
libgomp1 && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
# Set the entry point to activate the virtual environment and run the command line tool
# ENTRYPOINT ["/bin/bash", "-c", "source /opt/mineru_venv/bin/activate && exec \"$@\" && python3 app.py", "--"]
# Create volume for paddleocr models
RUN mkdir -p /root/.paddleocr
VOLUME [ "/root/.paddleocr" ]
# Copy the app and its configuration file
COPY entrypoint.sh /app/entrypoint.sh
COPY magic-pdf.json /root/magic-pdf.json
COPY app.py /app/app.py
# Expose the port that FastAPI will run on
EXPOSE 8000
# Command to run FastAPI using Uvicorn, pointing to app.py and binding to 0.0.0.0:8000
CMD ["/bin/bash", "-c", "source /opt/mineru_venv/bin/activate && uvicorn app:app --host 0.0.0.0 --port 8000"]
\ No newline at end of file
ENTRYPOINT [ "/app/entrypoint.sh" ]
CMD ["--host", "0.0.0.0", "--port", "8000"]
projects/web_api/README.md
View file @ 1f497129
基于MinerU的PDF解析API
# 基于MinerU的PDF解析API
- MinerU的GPU镜像构建
- 基于FastAPI的PDF解析接口
- MinerU的GPU镜像构建
- 基于FastAPI的PDF解析接口
支持一键启动，已经打包到镜像中，自带模型权重，支持GPU推理加速，GPU速度相比CPU每页解析要快几十倍不等
## 构建方式
```
docker build -t mineru-api .
```
## 启动命令：
或者使用代理：
```
docker build --build-arg http_proxy=http://127.0.0.1:7890 --build-arg https_proxy=http://127.0.0.1:7890 -t mineru-api .
```
```docker run -itd --name=mineru_server --gpus=all -p 8888:8000 quincyqiang/mineru:0.1-models```
## 启动命令
![](https://i-blog.csdnimg.cn/direct/bcff4f524ea5400db14421ba7cec4989.png)
```
docker run --rm -it --gpus=all -v ./paddleocr:/root/.paddleocr -p 8000:8000 mineru-api
```
具体截图请见博客：https://blog.csdn.net/yanqianglifei/article/details/141979684
初次调用 API 时会自动下载 paddleocr 的模型（约数十 MB），其余模型已包含在镜像中。
## 测试参数
## 启动日志：
访问地址：
![](https://i-blog.csdnimg.cn/direct/4eb5657567e4415eba912179dca5c8aa.png)
```
http://localhost:8000/docs
http://127.0.0.1:8000/docs
```
## 输入参数：
## 旧版镜像地址
访问地址：
> 阿里云地址：docker pull registry.cn-beijing.aliyuncs.com/quincyqiang/mineru:0.1-models
>
> dockerhub地址：docker pull quincyqiang/mineru:0.1-models
http://localhost:8888/docs
http://127.0.01:8888/docs
## 旧版截图
![](https://i-blog.csdnimg.cn/direct/8b3a2bc5908042268e8cc69756e331a2.png)
### 启动命令
## 解析效果：
![](https://i-blog.csdnimg.cn/direct/bcff4f524ea5400db14421ba7cec4989.png)
![](https://i-blog.csdnimg.cn/direct/a54dcae834ae48d498fb595aca4212c3.png)
具体截图请见博客：https://blog.csdn.net/yanqianglifei/article/details/141979684
### 启动日志
![](https://i-blog.csdnimg.cn/direct/4eb5657567e4415eba912179dca5c8aa.png)
## 镜像地址：
### 测试参数
> 阿里云地址：docker pull registry.cn-beijing.aliyuncs.com/quincyqiang/mineru:0.1-models
![](https://i-blog.csdnimg.cn/direct/8b3a2bc5908042268e8cc69756e331a2.png)
> dockerhub地址：docker pull quincyqiang/mineru:0.1-models
### 解析效果
![](https://i-blog.csdnimg.cn/direct/a54dcae834ae48d498fb595aca4212c3.png)
projects/web_api/app.py
View file @ 1f497129
import copy
import json
import os
from tempfile import NamedTemporaryFile
from base64 import b64encode
from glob import glob
from io import StringIO
from typing import Tuple, Union
import uvicorn
from fastapi import FastAPI, File, UploadFile
from fastapi import FastAPI, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from loguru import logger
import magic_pdf.model as model_config
from magic_pdf.config.enums import SupportedPdfParseMethod
from magic_pdf.data.data_reader_writer import FileBasedDataWriter
from magic_pdf.data.data_reader_writer import DataWriter, FileBasedDataWriter
from magic_pdf.data.data_reader_writer.s3 import S3DataReader, S3DataWriter
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.libs.config_reader import get_bucket_name, get_s3_config
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from magic_pdf.operators.models import InferenceResult
from magic_pdf.operators.pipes import PipeResult
model_config.__use_inside_model__ = True
app = FastAPI()
def json_md_dump(
model_json,
middle_json,
md_writer,
pdf_name,
content_list,
md_content,
):
# Write model results to model.json
orig_model_list = copy.deepcopy(model_json)
md_writer.write_string(
f'{pdf_name}_model.json',
json.dumps(orig_model_list, ensure_ascii=False, indent=4),
)
# Write intermediate results to middle.json
md_writer.write_string(
f'{pdf_name}_middle.json',
json.dumps(middle_json, ensure_ascii=False, indent=4),
)
# Write text content results to content_list.json
md_writer.write_string(
f'{pdf_name}_content_list.json',
json.dumps(content_list, ensure_ascii=False, indent=4),
)
# Write results to .md file
md_writer.write_string(
f'{pdf_name}.md',
md_content,
)
@app.post('/pdf_parse', tags=['projects'], summary='Parse PDF file')
async def pdf_parse_main(
pdf_file: UploadFile = File(...),
parse_method: str = 'auto',
model_json_path: str = None,
is_json_md_dump: bool = True,
output_dir: str = 'output',
):
"""Execute the process of converting PDF to JSON and MD, outputting MD and
JSON files to the specified directory.
:param pdf_file: The PDF file to be parsed
:param parse_method: Parsing method, can be auto, ocr, or txt. Default is auto. If results are not satisfactory, try ocr
:param model_json_path: Path to existing model data file. If empty, use built-in model. PDF and model_json must correspond
:param is_json_md_dump: Whether to write parsed data to .json and .md files. Default is True. Different stages of data will be written to different .json files (3 in total), md content will be saved to .md file # noqa E501
:param output_dir: Output directory for results. A folder named after the PDF file will be created to store all results
class MemoryDataWriter(DataWriter):
def __init__(self):
self.buffer = StringIO()
def write(self, path: str, data: bytes) -> None:
if isinstance(data, str):
self.buffer.write(data)
else:
self.buffer.write(data.decode("utf-8"))
def write_string(self, path: str, data: str) -> None:
self.buffer.write(data)
def get_value(self) -> str:
return self.buffer.getvalue()
def close(self):
self.buffer.close()
def init_writers(
pdf_path: str = None,
pdf_file: UploadFile = None,
output_path: str = None,
output_image_path: str = None,
) -> Tuple[
Union[S3DataWriter, FileBasedDataWriter],
Union[S3DataWriter, FileBasedDataWriter],
bytes,
]:
"""
try:
# Create a temporary file to store the uploaded PDF
with NamedTemporaryFile(delete=False, suffix='.pdf') as temp_pdf:
temp_pdf.write(await pdf_file.read())
temp_pdf_path = temp_pdf.name
Initialize writers based on path type
pdf_name = os.path.basename(pdf_file.filename).split('.')[0]
Args:
pdf_path: PDF file path (local path or S3 path)
pdf_file: Uploaded PDF file object
output_path: Output directory path
output_image_path: Image output directory path
if output_dir:
output_path = os.path.join(output_dir, pdf_name)
Returns:
Tuple[writer, image_writer, pdf_bytes]: Returns initialized writer tuple and PDF
file content
"""
if pdf_path:
is_s3_path = pdf_path.startswith("s3://")
if is_s3_path:
bucket = get_bucket_name(pdf_path)
ak, sk, endpoint = get_s3_config(bucket)
writer = S3DataWriter(
output_path, bucket=bucket, ak=ak, sk=sk, endpoint_url=endpoint
)
image_writer = S3DataWriter(
output_image_path, bucket=bucket, ak=ak, sk=sk, endpoint_url=endpoint
)
# 临时创建reader读取文件内容
temp_reader = S3DataReader(
"", bucket=bucket, ak=ak, sk=sk, endpoint_url=endpoint
)
pdf_bytes = temp_reader.read(pdf_path)
else:
output_path = os.path.join(os.path.dirname(temp_pdf_path), pdf_name)
writer = FileBasedDataWriter(output_path)
image_writer = FileBasedDataWriter(output_image_path)
os.makedirs(output_image_path, exist_ok=True)
with open(pdf_path, "rb") as f:
pdf_bytes = f.read()
else:
# 处理上传的文件
pdf_bytes = pdf_file.file.read()
writer = FileBasedDataWriter(output_path)
image_writer = FileBasedDataWriter(output_image_path)
os.makedirs(output_image_path, exist_ok=True)
output_image_path = os.path.join(output_path, 'images')
return writer, image_writer, pdf_bytes
# Get parent path of images for relative path in .md and content_list.json
image_path_parent = os.path.basename(output_image_path)
pdf_bytes = open(temp_pdf_path, 'rb').read() # Read binary data of PDF file
def process_pdf(
pdf_bytes: bytes,
parse_method: str,
image_writer: Union[S3DataWriter, FileBasedDataWriter],
) -> Tuple[InferenceResult, PipeResult]:
"""
Process PDF file content
Args:
pdf_bytes: Binary content of PDF file
parse_method: Parse method ('ocr', 'txt', 'auto')
image_writer: Image writer
if model_json_path:
# Read original JSON data of PDF file parsed by model, list type
model_json = json.loads(open(model_json_path, 'r', encoding='utf-8').read())
Returns:
Tuple[InferenceResult, PipeResult]: Returns inference result and pipeline result
"""
ds = PymuDocDataset(pdf_bytes)
infer_result: InferenceResult = None
pipe_result: PipeResult = None
if parse_method == "ocr":
infer_result = ds.apply(doc_analyze, ocr=True)
pipe_result = infer_result.pipe_ocr_mode(image_writer)
elif parse_method == "txt":
infer_result = ds.apply(doc_analyze, ocr=False)
pipe_result = infer_result.pipe_txt_mode(image_writer)
else: # auto
if ds.classify() == SupportedPdfParseMethod.OCR:
infer_result = ds.apply(doc_analyze, ocr=True)
pipe_result = infer_result.pipe_ocr_mode(image_writer)
else:
model_json = []
# Execute parsing steps
image_writer, md_writer = FileBasedDataWriter(
output_image_path
), FileBasedDataWriter(output_path)
ds = PymuDocDataset(pdf_bytes)
# Choose parsing method
if parse_method == 'auto':
if ds.classify() == SupportedPdfParseMethod.OCR:
parse_method = 'ocr'
else:
parse_method = 'txt'
if parse_method not in ['txt', 'ocr']:
logger.error('Unknown parse method, only auto, ocr, txt allowed')
infer_result = ds.apply(doc_analyze, ocr=False)
pipe_result = infer_result.pipe_txt_mode(image_writer)
return infer_result, pipe_result
def encode_image(image_path: str) -> str:
"""Encode image using base64"""
with open(image_path, "rb") as f:
return b64encode(f.read()).decode()
@app.post(
"/pdf_parse",
tags=["projects"],
summary="Parse PDF files (supports local files and S3)",
)
async def pdf_parse(
pdf_file: UploadFile = None,
pdf_path: str = None,
parse_method: str = "auto",
is_json_md_dump: bool = False,
output_dir: str = "output",
return_layout: bool = False,
return_info: bool = False,
return_content_list: bool = False,
return_images: bool = False,
):
"""
Execute the process of converting PDF to JSON and MD, outputting MD and JSON files
to the specified directory.
Args:
pdf_file: The PDF file to be parsed. Must not be specified together with
`pdf_path`
pdf_path: The path to the PDF file to be parsed. Must not be specified together
with `pdf_file`
parse_method: Parsing method, can be auto, ocr, or txt. Default is auto. If
results are not satisfactory, try ocr
is_json_md_dump: Whether to write parsed data to .json and .md files. Default
to False. Different stages of data will be written to different .json files
(3 in total), md content will be saved to .md file
output_dir: Output directory for results. A folder named after the PDF file
will be created to store all results
return_layout: Whether to return parsed PDF layout. Default to False
return_info: Whether to return parsed PDF info. Default to False
return_content_list: Whether to return parsed PDF content list. Default to False
"""
try:
if (pdf_file is None and pdf_path is None) or (
pdf_file is not None and pdf_path is not None
):
return JSONResponse(
content={'error': 'Invalid parse method'}, status_code=400
content={"error": "Must provide either pdf_file or pdf_path"},
status_code=400,
)
if len(model_json) == 0:
if parse_method == 'ocr':
infer_result = ds.apply(doc_analyze, ocr=True)
else:
infer_result = ds.apply(doc_analyze, ocr=False)
# Get PDF filename
pdf_name = os.path.basename(pdf_path if pdf_path else pdf_file.filename).split(
"."
)[0]
output_path = f"{output_dir}/{pdf_name}"
output_image_path = f"{output_path}/images"
else:
infer_result = InferenceResult(model_json, ds)
if len(model_json) == 0 and not model_config.__use_inside_model__:
logger.error('Need model list input')
return JSONResponse(
content={'error': 'Model list input required'}, status_code=400
)
if parse_method == 'ocr':
pipe_res = infer_result.pipe_ocr_mode(image_writer)
else:
pipe_res = infer_result.pipe_txt_mode(image_writer)
# Initialize readers/writers and get PDF content
writer, image_writer, pdf_bytes = init_writers(
pdf_path=pdf_path,
pdf_file=pdf_file,
output_path=output_path,
output_image_path=output_image_path,
)
# Process PDF
infer_result, pipe_result = process_pdf(pdf_bytes, parse_method, image_writer)
# Use MemoryDataWriter to get results
content_list_writer = MemoryDataWriter()
md_content_writer = MemoryDataWriter()
middle_json_writer = MemoryDataWriter()
# Use PipeResult's dump method to get data
pipe_result.dump_content_list(content_list_writer, "", "images")
pipe_result.dump_md(md_content_writer, "", "images")
pipe_result.dump_middle_json(middle_json_writer, "")
# Save results in text and md format
content_list = pipe_res.get_content_list(image_path_parent, drop_mode='none')
md_content = pipe_res.get_markdown(image_path_parent, drop_mode='none')
# Get content
content_list = json.loads(content_list_writer.get_value())
md_content = md_content_writer.get_value()
middle_json = json.loads(middle_json_writer.get_value())
model_json = infer_result.get_infer_res()
# If results need to be saved
if is_json_md_dump:
json_md_dump(infer_result._infer_res, pipe_res._pipe_res, md_writer, pdf_name, content_list, md_content)
data = {
'layout': copy.deepcopy(infer_result._infer_res),
'info': pipe_res._pipe_res,
'content_list': content_list,
'md_content': md_content,
}
writer.write_string(
f"{pdf_name}_content_list.json", content_list_writer.get_value()
)
writer.write_string(f"{pdf_name}.md", md_content)
writer.write_string(
f"{pdf_name}_middle.json", middle_json_writer.get_value()
)
writer.write_string(
f"{pdf_name}_model.json",
json.dumps(model_json, indent=4, ensure_ascii=False),
)
# Save visualization results
pipe_result.draw_layout(os.path.join(output_path, f"{pdf_name}_layout.pdf"))
pipe_result.draw_span(os.path.join(output_path, f"{pdf_name}_spans.pdf"))
pipe_result.draw_line_sort(
os.path.join(output_path, f"{pdf_name}_line_sort.pdf")
)
infer_result.draw_model(os.path.join(output_path, f"{pdf_name}_model.pdf"))
# Build return data
data = {}
if return_layout:
data["layout"] = model_json
if return_info:
data["info"] = middle_json
if return_content_list:
data["content_list"] = content_list
if return_images:
image_paths = glob(f"{output_image_path}/*.jpg")
data["images"] = {
os.path.basename(
image_path
): f"data:image/jpeg;base64,{encode_image(image_path)}"
for image_path in image_paths
}
data["md_content"] = md_content # md_content is always returned
# Clean up memory writers
content_list_writer.close()
md_content_writer.close()
middle_json_writer.close()
return JSONResponse(data, status_code=200)
except Exception as e:
logger.exception(e)
return JSONResponse(content={'error': str(e)}, status_code=500)
finally:
# Clean up the temporary file
if 'temp_pdf_path' in locals():
os.unlink(temp_pdf_path)
return JSONResponse(content={"error": str(e)}, status_code=500)
if __name__ == '__main__':
uvicorn.run(app, host='0.0.0.0', port=8888)
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8888)
projects/web_api/download_models.py 0 → 100755
View file @ 1f497129
#!/usr/bin/env python
from huggingface_hub import snapshot_download
if __name__ == "__main__":
mineru_patterns = [
"models/Layout/LayoutLMv3/*",
"models/Layout/YOLO/*",
"models/MFD/YOLO/*",
"models/MFR/unimernet_small_2501/*",
"models/TabRec/TableMaster/*",
"models/TabRec/StructEqTable/*",
]
model_dir = snapshot_download(
"opendatalab/PDF-Extract-Kit-1.0",
allow_patterns=mineru_patterns,
local_dir="/opt/",
)
layoutreader_pattern = [
"*.json",
"*.safetensors",
]
layoutreader_model_dir = snapshot_download(
"hantian/layoutreader",
allow_patterns=layoutreader_pattern,
local_dir="/opt/layoutreader/",
)
model_dir = model_dir + "/models"
print(f"model_dir is: {model_dir}")
print(f"layoutreader_model_dir is: {layoutreader_model_dir}")
projects/web_api/entrypoint.sh 0 → 100755
View file @ 1f497129
#!/usr/bin/env bash
set -euo pipefail
. /app/venv/bin/activate
exec uvicorn app:app "$@"
projects/web_api/magic-pdf.json
View file @ 1f497129
......@@ -4,10 +4,41 @@
"bucket-name-2":["ak", "sk", "endpoint"]
},
"models-dir":"/opt/models",
"layoutreader-model-dir":"/opt/layoutreader",
"device-mode":"cuda",
"layout-config": {
"model": "doclayout_yolo"
},
"formula-config": {
"mfd_model": "yolo_v8_mfd",
"mfr_model": "unimernet_small",
"enable": true
},
"table-config": {
"model": "TableMaster",
"is_table_recog_enable": false,
"model": "rapid_table",
"sub_model": "slanet_plus",
"enable": true,
"max_time": 400
}
},
"llm-aided-config": {
"formula_aided": {
"api_key": "your_api_key",
"base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
"model": "qwen2.5-7b-instruct",
"enable": false
},
"text_aided": {
"api_key": "your_api_key",
"base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
"model": "qwen2.5-7b-instruct",
"enable": false
},
"title_aided": {
"api_key": "your_api_key",
"base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1",
"model": "qwen2.5-32b-instruct",
"enable": false
}
},
"config_version": "1.1.1"
}
projects/web_api/magic-pdf.template.json deleted 100644 → 0
View file @ 9a87d3ea
{
"bucket_info":{
"bucket-name-1":["ak", "sk", "endpoint"],
"bucket-name-2":["ak", "sk", "endpoint"]
},
"models-dir":"/tmp/models",
"device-mode":"cuda",
"table-config": {
"model": "TableMaster",
"is_table_recog_enable": false,
"max_time": 400
}
}
\ No newline at end of file
projects/web_api/requirements.txt
View file @ 1f497129
--extra-index-url https://myhloli.github.io/wheels/
magic-pdf[full]
fastapi
uvicorn
python-multipart
projects/web_api/sources.list deleted 100644 → 0
View file @ 9a87d3ea
deb http://mirrors.aliyun.com/ubuntu/ focal main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ focal main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ focal-security main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ focal-security main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ focal-updates main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ focal-updates main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ focal-proposed main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ focal-proposed main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ focal-backports main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ focal-backports main restricted universe multiverse
\ No newline at end of file
projects/web_api/start_mineru.sh deleted 100644 → 0
View file @ 9a87d3ea
docker run -itd --name=mineru_server --gpus=all -p 8888:8000 quincyqiang/mineru:0.1-models /bin/bash
docker run -itd --name=mineru_server --gpus=all -p 8888:8000 quincyqiang/mineru:0.3-models
docker login --username=1185918903@qq.com registry.cn-beijing.aliyuncs.com
docker tag quincyqiang/mineru:0.3-models registry.cn-beijing.aliyuncs.com/quincyqiang/gomate:0.3-models
docker push registry.cn-beijing.aliyuncs.com/quincyqiang/gomate:0.3-models
\ No newline at end of file
projects/web_demo/web_demo/common/ext.py
View file @ 1f497129
import hashlib
import mimetypes
import urllib.parse
def is_pdf(filename, file):
"""
判断文件是否为PDF格式。
判断文件是否为PDF格式，支持中文名和特殊字符。
:param filename: 文件名
:param file: 文件对象
:return: 如果文件是PDF格式，则返回True，否则返回False
"""
# 检查文件扩展名 https://arxiv.org/pdf/2405.08702 pdf链接可能存在不带扩展名的情况，先注释
# if not filename.endswith('.pdf'):
# return False
# 检查MIME类型
mime_type, _ = mimetypes.guess_type(filename)
print(mime_type)
if mime_type != 'application/pdf':
return False
# 可选：读取文件的前几KB内容并检查MIME类型
# 这一步是可选的，用于更严格的检查
# if not mimetypes.guess_type(filename, strict=False)[0] == 'application/pdf':
# return False
# 检查文件内容
file_start = file.read(5)
file.seek(0)
if not file_start.startswith(b'%PDF-'):
return False
return True
try:
# 对文件名进行URL解码，处理特殊字符
decoded_filename = urllib.parse.unquote(filename)
# 检查MIME类型
mime_type, _ = mimetypes.guess_type(decoded_filename)
print(f"Detected MIME type: {mime_type}")
# 某些情况下mime_type可能为None，需要特殊处理
if mime_type is None:
# 只检查文件内容的PDF标识
file_start = file.read(5)
file.seek(0) # 重置文件指针
return file_start.startswith(b'%PDF-')
if mime_type != 'application/pdf':
return False
# 检查文件内容的PDF标识
file_start = file.read(5)
file.seek(0) # 重置文件指针
if not file_start.startswith(b'%PDF-'):
return False
return True
except Exception as e:
print(f"Error checking PDF format: {str(e)}")
# 发生错误时，仍然尝试通过文件头判断
try:
file_start = file.read(5)
file.seek(0)
return file_start.startswith(b'%PDF-')
except:
return False
def url_is_pdf(file):
......
setup.py
View file @ 1f497129
......@@ -43,14 +43,14 @@ if __name__ == '__main__':
"matplotlib;platform_system=='Linux' or platform_system=='Darwin'", # linux 和 macos 不应限制matplotlib的最高版本，以避免无法更新导致的一些bug
"ultralytics>=8.3.48", # yolov8,公式检测
"paddleocr==2.7.3", # 2.8.0及2.8.1版本与detectron2有冲突，需锁定2.7.3
"paddlepaddle==3.0.0b1;platform_system=='Linux'", # 解决linux的段异常问题
"paddlepaddle==2.6.1;platform_system=='Windows' or platform_system=='Darwin'", # windows版本3.0.0b1效率下降，需锁定2.6.1
"paddlepaddle==3.0.0rc1;platform_system=='Linux' or platform_system=='Darwin'", # 解决linux的段异常问题
"paddlepaddle==2.6.1;platform_system=='Windows'", # windows版本3.0.0效率下降，需锁定2.6.1
"struct-eqtable==0.3.2", # 表格解析
"einops", # struct-eqtable依赖
"accelerate", # struct-eqtable依赖
"doclayout_yolo==0.0.2b1", # doclayout_yolo
"rapidocr-paddle", # rapidocr-paddle
"rapidocr_onnxruntime",
"rapidocr-paddle>=1.4.5,<2.0.0", # rapidocr-paddle
"rapidocr_onnxruntime>=1.4.4,<2.0.0",
"rapid_table>=1.0.3,<2.0.0", # rapid_table
"PyYAML", # yaml
"openai", # openai SDK
......
tests/unittest/test_integrations/test_rag/test_utils.py
View file @ 1f497129
......@@ -24,7 +24,7 @@ def test_convert_middle_json_to_layout_elements():
assert len(res[0].layout_dets) > 0
assert res[0].layout_dets[0].anno_id == 0
assert res[0].layout_dets[0].category_type == CategoryType.text
assert len(res[0].extra.element_relation) >= 3
assert len(res[0].extra.element_relation) >= 2
# teardown
shutil.rmtree(temp_output_dir)
......@@ -51,7 +51,7 @@ def test_inference():
assert len(res[0].layout_dets) > 0
assert res[0].layout_dets[0].anno_id == 0
assert res[0].layout_dets[0].category_type == CategoryType.text
assert len(res[0].extra.element_relation) >= 3
assert len(res[0].extra.element_relation) >= 2
# teardown
shutil.rmtree(temp_output_dir)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment