Update app.py, count_pdfs.py, LICENSE.md, magic-pdf.template.json,...

Update app.py, count_pdfs.py, LICENSE.md, magic-pdf.template.json, requirements.txt, requirements-docker.txt, requirements-qa.txt, update_version.py, setup.py, magic_pdf/__init__.py, magic_pdf/pdf_parse_by_txt.py, magic_pdf/pdf_parse_by_ocr.py, magic_pdf/user_api.py, magic_pdf/pdf_parse_union_core.py, magic_pdf/__pycache__/pdf_parse_by_ocr.cpython-310.pyc, magic_pdf/__pycache__/__init__.cpython-310.pyc, magic_pdf/__pycache__/pdf_parse_by_txt.cpython-310.pyc, magic_pdf/__pycache__/pdf_parse_union_core.cpython-310.pyc, magic_pdf/__pycache__/user_api.cpython-310.pyc, magic_pdf/dict2md/__init__.py, magic_pdf/dict2md/mkcontent.py, magic_pdf/dict2md/ocr_mkcontent.py, magic_pdf/dict2md/ocr_client.py, magic_pdf/dict2md/ocr_server.py, magic_pdf/dict2md/ocr_server_72.py, magic_pdf/dict2md/tmp.py, magic_pdf/dict2md/__pycache__/__init__.cpython-310.pyc, magic_pdf/dict2md/__pycache__/ocr_client.cpython-310.pyc, magic_pdf/dict2md/__pycache__/ocr_mkcontent.cpython-310.pyc, magic_pdf/filter/__init__.py, magic_pdf/filter/pdf_classify_by_type.py, magic_pdf/filter/pdf_meta_scan.py, magic_pdf/integrations/__init__.py, magic_pdf/integrations/rag/__init__.py, magic_pdf/integrations/rag/api.py, magic_pdf/integrations/rag/type.py, magic_pdf/integrations/rag/utils.py, magic_pdf/layout/__init__.py, magic_pdf/layout/bbox_sort.py, magic_pdf/layout/layout_det_utils.py, magic_pdf/layout/layout_sort.py, magic_pdf/layout/layout_spiler_recog.py, magic_pdf/layout/mcol_sort.py, magic_pdf/libs/__init__.py, magic_pdf/libs/boxbase.py, magic_pdf/libs/calc_span_stats.py, magic_pdf/libs/commons.py, magic_pdf/libs/config_reader.py, magic_pdf/libs/Constants.py, magic_pdf/libs/convert_utils.py, magic_pdf/libs/coordinate_transform.py, magic_pdf/libs/detect_language_from_model.py, magic_pdf/libs/draw_bbox.py, magic_pdf/libs/drop_reason.py, magic_pdf/libs/drop_tag.py, magic_pdf/libs/hash_utils.py, magic_pdf/libs/json_compressor.py, magic_pdf/libs/language.py, magic_pdf/libs/local_math.py, magic_pdf/libs/MakeContentConfig.py, magic_pdf/libs/markdown_utils.py, magic_pdf/libs/nlp_utils.py, magic_pdf/libs/ModelBlockTypeEnum.py, magic_pdf/libs/ocr_content_type.py, magic_pdf/libs/path_utils.py, magic_pdf/libs/pdf_check.py, magic_pdf/libs/pdf_image_tools.py, magic_pdf/libs/safe_filename.py, magic_pdf/libs/textbase.py, magic_pdf/libs/version.py, magic_pdf/libs/vis_utils.py, magic_pdf/model/__init__.py, magic_pdf/model/doc_analyze_by_custom_model.py, magic_pdf/model/magic_model.py, magic_pdf/model/pdf_extract_kit.py, magic_pdf/model/model_list.py, magic_pdf/model/pp_structure_v2.py, magic_pdf/model/ppTableModel.py, magic_pdf/model/pek_sub_modules/__init__.py, magic_pdf/model/pek_sub_modules/post_process.py, magic_pdf/model/pek_sub_modules/self_modify.py, magic_pdf/model/pek_sub_modules/layoutlmv3/__init__.py, magic_pdf/model/pek_sub_modules/layoutlmv3/backbone.py, magic_pdf/model/pek_sub_modules/layoutlmv3/beit.py, magic_pdf/model/pek_sub_modules/layoutlmv3/deit.py, magic_pdf/model/pek_sub_modules/layoutlmv3/model_init.py, magic_pdf/model/pek_sub_modules/layoutlmv3/rcnn_vl.py, magic_pdf/model/pek_sub_modules/layoutlmv3/visualizer.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/__init__.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/models/__init__.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/models/layoutlmv3/__init__.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/models/layoutlmv3/modeling_layoutlmv3.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/models/layoutlmv3/tokenization_layoutlmv3_fast.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/models/layoutlmv3/configuration_layoutlmv3.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/models/layoutlmv3/tokenization_layoutlmv3.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/data/cord.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/data/__init__.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/data/data_collator.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/data/funsd.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/data/image_utils.py, magic_pdf/model/pek_sub_modules/layoutlmv3/layoutlmft/data/xfund.py, magic_pdf/model/pek_sub_modules/structeqtable/__init__.py, magic_pdf/model/pek_sub_modules/structeqtable/StructTableModel.py, magic_pdf/para/__init__.py, magic_pdf/para/block_continuation_processor.py, magic_pdf/para/block_termination_processor.py, magic_pdf/para/commons.py, magic_pdf/para/denoise.py, magic_pdf/para/draw.py, magic_pdf/para/exceptions.py, magic_pdf/para/layout_match_processor.py, magic_pdf/para/para_pipeline.py, magic_pdf/para/para_split.py, magic_pdf/para/para_split_v2.py, magic_pdf/para/raw_processor.py, magic_pdf/para/stats.py, magic_pdf/para/title_processor.py, magic_pdf/parse/__init__.py, magic_pdf/parse/common_parse.py, magic_pdf/parse/excel_parse.py, magic_pdf/parse/pdf_client.py, magic_pdf/pipe/__init__.py, magic_pdf/pipe/AbsPipe.py, magic_pdf/pipe/OCRPipe.py, magic_pdf/pipe/TXTPipe.py, magic_pdf/pipe/UNIPipe.py, magic_pdf/post_proc/__init__.py, magic_pdf/post_proc/pdf_post_filter.py, magic_pdf/post_proc/remove_footnote.py, magic_pdf/post_proc/detect_para.py, magic_pdf/pre_proc/__init__.py, magic_pdf/pre_proc/citationmarker_remove.py, magic_pdf/pre_proc/construct_page_dict.py, magic_pdf/pre_proc/cut_image.py, magic_pdf/pre_proc/detect_equation.py, magic_pdf/pre_proc/detect_footer_header_by_statistics.py, magic_pdf/pre_proc/detect_footer_by_model.py, magic_pdf/pre_proc/detect_footnote.py, magic_pdf/pre_proc/detect_header.py, magic_pdf/pre_proc/detect_images.py, magic_pdf/pre_proc/detect_page_number.py, magic_pdf/pre_proc/detect_tables.py, magic_pdf/pre_proc/equations_replace.py, magic_pdf/pre_proc/fix_image.py, magic_pdf/pre_proc/fix_table.py, magic_pdf/pre_proc/main_text_font.py, magic_pdf/pre_proc/ocr_detect_all_bboxes.py, magic_pdf/pre_proc/ocr_detect_layout.py, magic_pdf/pre_proc/ocr_dict_merge.py, magic_pdf/pre_proc/ocr_span_list_modify.py, magic_pdf/pre_proc/post_layout_split.py, magic_pdf/pre_proc/remove_bbox_overlap.py, magic_pdf/pre_proc/remove_colored_strip_bbox.py, magic_pdf/pre_proc/pdf_pre_filter.py, magic_pdf/pre_proc/remove_footer_header.py, magic_pdf/pre_proc/remove_rotate_bbox.py, magic_pdf/pre_proc/resolve_bbox_conflict.py, magic_pdf/pre_proc/solve_line_alien.py, magic_pdf/pre_proc/statistics.py, magic_pdf/resources/fasttext-langdetect/lid.176.ftz, magic_pdf/resources/model_config/model_configs.yaml, magic_pdf/resources/model_config/layoutlmv3/layoutlmv3_base_inference.yaml, magic_pdf/resources/model_config/UniMERNet/demo.yaml, magic_pdf/rw/__init__.py, magic_pdf/rw/AbsReaderWriter.py, magic_pdf/rw/DiskReaderWriter.py, magic_pdf/rw/S3ReaderWriter.py, magic_pdf/spark/__init__.py, magic_pdf/spark/spark_api.py, magic_pdf/tools/__init__.py, magic_pdf/tools/pdf_client.py, magic_pdf/tools/common.py, magic_pdf/tools/cli_dev.py, magic_pdf/tools/cli.py, magic_pdf/tools/pdf_server.py files

app.py

+# Copyright (c) Opendatalab. All rights reserved.
+
+import base64
+import os
+import time
+import zipfile
+from pathlib import Path
+import re
+
+from loguru import logger
+
+from magic_pdf.libs.hash_utils import compute_sha256
+from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
+from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
+from magic_pdf.tools.common import do_parse, prepare_env
+
+os.system("pip install gradio")
+os.system("pip install gradio-pdf")
+import gradio as gr
+from gradio_pdf import PDF
+
+
+def read_fn(path):
+    disk_rw = DiskReaderWriter(os.path.dirname(path))
+    return disk_rw.read(os.path.basename(path), AbsReaderWriter.MODE_BIN)
+
+
+def parse_pdf(doc_path, output_dir, end_page_id):
+    os.makedirs(output_dir, exist_ok=True)
+
+    try:
+        file_name = f"{str(Path(doc_path).stem)}_{time.time()}"
+        pdf_data = read_fn(doc_path)
+        parse_method = "auto"
+        local_image_dir, local_md_dir = prepare_env(output_dir, file_name, parse_method)
+        do_parse(
+            output_dir,
+            file_name,
+            pdf_data,
+            [],
+            parse_method,
+            False,
+            end_page_id=end_page_id,
+        )
+        return local_md_dir, file_name
+    except Exception as e:
+        logger.exception(e)
+
+
+def compress_directory_to_zip(directory_path, output_zip_path):
+    """
+    压缩指定目录到一个 ZIP 文件。
+
+    :param directory_path: 要压缩的目录路径
+    :param output_zip_path: 输出的 ZIP 文件路径
+    """
+    try:
+        with zipfile.ZipFile(output_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+
+            # 遍历目录中的所有文件和子目录
+            for root, dirs, files in os.walk(directory_path):
+                for file in files:
+                    # 构建完整的文件路径
+                    file_path = os.path.join(root, file)
+                    # 计算相对路径
+                    arcname = os.path.relpath(file_path, directory_path)
+                    # 添加文件到 ZIP 文件
+                    zipf.write(file_path, arcname)
+        return 0
+    except Exception as e:
+        logger.exception(e)
+        return -1
+
+
+def image_to_base64(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode('utf-8')
+
+
+def replace_image_with_base64(markdown_text, image_dir_path):
+    # 匹配Markdown中的图片标签
+    pattern = r'\!\[(?:[^\]]*)\]\(([^)]+)\)'
+
+    # 替换图片链接
+    def replace(match):
+        relative_path = match.group(1)
+        full_path = os.path.join(image_dir_path, relative_path)
+        base64_image = image_to_base64(full_path)
+        return f"![{relative_path}](data:image/jpeg;base64,{base64_image})"
+
+    # 应用替换
+    return re.sub(pattern, replace, markdown_text)
+
+
+def to_markdown(file_path, end_pages):
+    # 获取识别的md文件以及压缩包文件路径
+    local_md_dir, file_name = parse_pdf(file_path, './output', end_pages - 1)
+    archive_zip_path = os.path.join("./output", compute_sha256(local_md_dir) + ".zip")
+    zip_archive_success = compress_directory_to_zip(local_md_dir, archive_zip_path)
+    if zip_archive_success == 0:
+        logger.info("压缩成功")
+    else:
+        logger.error("压缩失败")
+    md_path = os.path.join(local_md_dir, file_name + ".md")
+    with open(md_path, 'r', encoding='utf-8') as f:
+        txt_content = f.read()
+    md_content = replace_image_with_base64(txt_content, local_md_dir)
+    # 返回转换后的PDF路径
+    new_pdf_path = os.path.join(local_md_dir, file_name + "_layout.pdf")
+
+    return md_content, txt_content, archive_zip_path, new_pdf_path
+
+
+# def show_pdf(file_path):
+#     with open(file_path, "rb") as f:
+#         base64_pdf = base64.b64encode(f.read()).decode('utf-8')
+#     pdf_display = f'<embed src="data:application/pdf;base64,{base64_pdf}" ' \
+#                   f'width="100%" height="1000" type="application/pdf">'
+#     return pdf_display
+
+
+latex_delimiters = [{"left": "$$", "right": "$$", "display": True},
+                    {"left": '$', "right": '$', "display": False}]
+
+
+def init_model():
+    from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
+    try:
+        model_manager = ModelSingleton()
+        txt_model = model_manager.get_model(False, False)
+        logger.info(f"txt_model init final")
+        ocr_model = model_manager.get_model(True, False)
+        logger.info(f"ocr_model init final")
+        return 0
+    except Exception as e:
+        logger.exception(e)
+        return -1
+
+
+model_init = init_model()
+logger.info(f"model_init: {model_init}")
+
+
+if __name__ == "__main__":
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column(variant='panel', scale=5):
+                pdf_show = gr.Markdown()
+                max_pages = gr.Slider(1, 10, 5, step=1, label="Max convert pages")
+                with gr.Row() as bu_flow:
+                    change_bu = gr.Button("Convert")
+                    clear_bu = gr.ClearButton([pdf_show], value="Clear")
+                pdf_show = PDF(label="Please upload pdf", interactive=True, height=800)
+
+            with gr.Column(variant='panel', scale=5):
+                output_file = gr.File(label="convert result", interactive=False)
+                with gr.Tabs():
+                    with gr.Tab("Markdown rendering"):
+                        md = gr.Markdown(label="Markdown rendering", height=900, show_copy_button=True,
+                                         latex_delimiters=latex_delimiters, line_breaks=True)
+                    with gr.Tab("Markdown text"):
+                        md_text = gr.TextArea(lines=45, show_copy_button=True)
+        change_bu.click(fn=to_markdown, inputs=[pdf_show, max_pages], outputs=[md, md_text, output_file, pdf_show])
+        clear_bu.add([md, pdf_show, md_text, output_file])
+
+    demo.launch()
+

count_pdfs.py

+import os
+
+
+def find_pdfs_and_save_to_txt(directory, output_file):
+    pdf_paths = []
+
+    # Traverse the directory and find all PDF files
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            if file.lower().endswith('.pdf'):
+                pdf_paths.append(os.path.join(root, file))
+
+    # Write all PDF paths to the output file
+    with open(output_file, 'w',encoding='utf-8') as f:
+        for path in pdf_paths:
+            f.write(path + '\n')
+
+    return len(pdf_paths)
+
+#
+# # Use '.' for the current directory, and specify output file
+directory = '/home/practice/rag/DCU知识库文档集合'
+output_file = 'pdf_paths.txt'
+pdf_count = find_pdfs_and_save_to_txt(directory, output_file)
+#
+print(f"Found {pdf_count} PDF files. Paths are saved in '{output_file}'.")
+
+
+
+def remove_duplicate_lines(file_path):
+    lines_seen = set()  # 用于保存已经看到的唯一行
+    unique_lines = []  # 用于保存不重复的行
+
+    # 读取文件，并过滤重复的行
+    with open(file_path, 'r',encoding='utf-8') as f:
+        for line in f:
+            line = line.strip()  # 去掉行首尾的空白符
+            if line not in lines_seen:
+                unique_lines.append(line)
+                lines_seen.add(line)
+    out_path = file_path.split('.txt')[0] + '_uq.txt'
+    # 将不重复的行写回文件
+    with open(out_path, 'w',encoding='utf-8') as f:
+        for line in unique_lines:
+            f.write(line + '\n')
+
+    # 返回不重复行的数量
+    return len(unique_lines)
+
+
+# 指定要处理的文件路径
+file_path = 'pdf_paths.txt'
+unique_line_count = remove_duplicate_lines(file_path)
+
+print(f"Number of unique lines: {unique_line_count}")
+
+
+

magic-pdf.template.json

+{
+    "bucket_info":{
+        "bucket-name-1":["ak", "sk", "endpoint"],
+        "bucket-name-2":["ak", "sk", "endpoint"]
+    },
+    "models-dir":"/home/practice/model/PDF-Extract-Kit/models",
+    "device-mode":"cuda",
+    "table-config": {
+        "model": "struct_eqtable",
+        "is_table_recog_enable": True,
+        "max_time": 400
+    }
+}

magic_pdf/dict2md/mkcontent.py

+import math
+from loguru import logger
+
+from magic_pdf.libs.boxbase import find_bottom_nearest_text_bbox, find_top_nearest_text_bbox
+from magic_pdf.libs.commons import join_path
+from magic_pdf.libs.ocr_content_type import ContentType
+
+TYPE_INLINE_EQUATION = ContentType.InlineEquation
+TYPE_INTERLINE_EQUATION = ContentType.InterlineEquation
+UNI_FORMAT_TEXT_TYPE = ['text', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']
+
+
+@DeprecationWarning
+def mk_nlp_markdown_1(para_dict: dict):
+    """
+    对排序后的bboxes拼接内容
+    """
+    content_lst = []
+    for _, page_info in para_dict.items():
+        para_blocks = page_info.get("para_blocks")
+        if not para_blocks:
+            continue
+
+        for block in para_blocks:
+            item = block["paras"]
+            for _, p in item.items():
+                para_text = p["para_text"]
+                is_title = p["is_para_title"]
+                title_level = p['para_title_level']
+                md_title_prefix = "#"*title_level
+                if is_title:
+                    content_lst.append(f"{md_title_prefix} {para_text}")
+                else:
+                    content_lst.append(para_text)
+
+    content_text = "\n\n".join(content_lst)
+
+    return content_text
+
+
+
+# 找到目标字符串在段落中的索引
+def __find_index(paragraph, target):
+    index = paragraph.find(target)
+    if index != -1:
+        return index
+    else:
+        return None
+
+
+def __insert_string(paragraph, target, postion):
+    new_paragraph = paragraph[:postion] + target + paragraph[postion:] 
+    return new_paragraph
+
+
+def __insert_after(content, image_content, target):
+    """
+    在content中找到target，将image_content插入到target后面
+    """
+    index = content.find(target)
+    if index != -1:
+        content = content[:index+len(target)] + "\n\n" + image_content + "\n\n" + content[index+len(target):]
+    else:
+        logger.error(f"Can't find the location of image {image_content} in the markdown file, search target is {target}")
+    return content
+
+def __insert_before(content, image_content, target):
+    """
+    在content中找到target，将image_content插入到target前面
+    """
+    index = content.find(target)
+    if index != -1:
+        content = content[:index] + "\n\n" + image_content + "\n\n" + content[index:]
+    else:
+        logger.error(f"Can't find the location of image {image_content} in the markdown file, search target is {target}")
+    return content
+
+
+@DeprecationWarning
+def mk_mm_markdown_1(para_dict: dict):
+    """拼装多模态markdown"""
+    content_lst = []
+    for _, page_info in para_dict.items():
+        page_lst = [] # 一个page内的段落列表
+        para_blocks = page_info.get("para_blocks")
+        pymu_raw_blocks = page_info.get("preproc_blocks")
+        
+        all_page_images = []
+        all_page_images.extend(page_info.get("images",[]))
+        all_page_images.extend(page_info.get("image_backup", []) )
+        all_page_images.extend(page_info.get("tables",[]))
+        all_page_images.extend(page_info.get("table_backup",[]) )
+        
+        if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景
+            for img in all_page_images:
+                page_lst.append(f"![]({img['image_path']})") # TODO 图片顺序
+            page_md = "\n\n".join(page_lst)
+            
+        else:
+            for block in para_blocks:
+                item = block["paras"]
+                for _, p in item.items():
+                    para_text = p["para_text"]
+                    is_title = p["is_para_title"]
+                    title_level = p['para_title_level']
+                    md_title_prefix = "#"*title_level
+                    if is_title:
+                        page_lst.append(f"{md_title_prefix} {para_text}")
+                    else:
+                        page_lst.append(para_text)
+                        
+            """拼装成一个页面的文本"""
+            page_md = "\n\n".join(page_lst)
+            """插入图片"""
+            for img in all_page_images:
+                imgbox = img['bbox']
+                img_content = f"![]({img['image_path']})"
+                # 先看在哪个block内
+                for block in pymu_raw_blocks:
+                    bbox = block['bbox']
+                    if bbox[0]-1 <= imgbox[0] < bbox[2]+1 and bbox[1]-1 <= imgbox[1] < bbox[3]+1:# 确定在block内
+                        for l in block['lines']:
+                            line_box = l['bbox']
+                            if line_box[0]-1 <= imgbox[0] < line_box[2]+1 and line_box[1]-1 <= imgbox[1] < line_box[3]+1: # 在line内的，插入line前面
+                                line_txt = "".join([s['text'] for s in l['spans']])
+                                page_md = __insert_before(page_md, img_content, line_txt)
+                                break
+                            break
+                        else:# 在行与行之间
+                            # 找到图片x0,y0与line的x0,y0最近的line
+                            min_distance = 100000
+                            min_line = None
+                            for l in block['lines']:
+                                line_box = l['bbox']
+                                distance = math.sqrt((line_box[0] - imgbox[0])**2 + (line_box[1] - imgbox[1])**2)
+                                if distance < min_distance:
+                                    min_distance = distance
+                                    min_line = l
+                            if min_line:
+                                line_txt = "".join([s['text'] for s in min_line['spans']])
+                                img_h = imgbox[3] - imgbox[1]
+                                if min_distance<img_h: # 文字在图片前面
+                                    page_md = __insert_after(page_md, img_content, line_txt)
+                                else:
+                                    page_md = __insert_before(page_md, img_content, line_txt)
+                            else:
+                                logger.error(f"Can't find the location of image {img['image_path']} in the markdown file #1")
+                else:# 应当在两个block之间
+                    # 找到上方最近的block，如果上方没有就找大下方最近的block
+                    top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, imgbox)
+                    if top_txt_block:
+                        line_txt = "".join([s['text'] for s in top_txt_block['lines'][-1]['spans']])
+                        page_md = __insert_after(page_md, img_content, line_txt)
+                    else:
+                        bottom_txt_block = find_bottom_nearest_text_bbox(pymu_raw_blocks, imgbox)
+                        if bottom_txt_block:
+                            line_txt = "".join([s['text'] for s in bottom_txt_block['lines'][0]['spans']])
+                            page_md = __insert_before(page_md, img_content, line_txt)
+                        else:
+                            logger.error(f"Can't find the location of image {img['image_path']} in the markdown file #2")
+                    
+        content_lst.append(page_md)
+                    
+    """拼装成全部页面的文本"""
+    content_text = "\n\n".join(content_lst)
+
+    return content_text
+
+
+def __insert_after_para(text, type, element, content_list):
+    """
+    在content_list中找到text，将image_path作为一个新的node插入到text后面
+    """
+    for i, c in enumerate(content_list):
+        content_type = c.get("type")
+        if content_type in UNI_FORMAT_TEXT_TYPE and text in c.get("text", ''):
+            if type == "image":
+                content_node = {
+                    "type": "image",
+                    "img_path": element.get("image_path"),
+                    "img_alt": "",
+                    "img_title": "",
+                    "img_caption": "",
+                }
+            elif type == "table":
+                content_node = {
+                    "type": "table",
+                    "img_path": element.get("image_path"),
+                    "table_latex": element.get("text"),
+                    "table_title": "",
+                    "table_caption": "",
+                    "table_quality": element.get("quality"),
+                }
+            content_list.insert(i+1, content_node)
+            break
+    else:
+        logger.error(f"Can't find the location of image {element.get('image_path')} in the markdown file, search target is {text}")
+    
+
+
+def __insert_before_para(text, type, element, content_list):
+    """
+    在content_list中找到text，将image_path作为一个新的node插入到text前面
+    """
+    for i, c in enumerate(content_list):
+        content_type = c.get("type")
+        if content_type in  UNI_FORMAT_TEXT_TYPE and text in c.get("text", ''):
+            if type == "image":
+                content_node = {
+                    "type": "image",
+                    "img_path": element.get("image_path"),
+                    "img_alt": "",
+                    "img_title": "",
+                    "img_caption": "",
+                }
+            elif type == "table":
+                content_node = {
+                    "type": "table",
+                    "img_path": element.get("image_path"),
+                    "table_latex": element.get("text"),
+                    "table_title": "",
+                    "table_caption": "",
+                    "table_quality": element.get("quality"),
+                }
+            content_list.insert(i, content_node)
+            break
+    else:
+        logger.error(f"Can't find the location of image {element.get('image_path')} in the markdown file, search target is {text}")
+         
+
+def mk_universal_format(pdf_info_list: list, img_buket_path):
+    """
+    构造统一格式 https://aicarrier.feishu.cn/wiki/FqmMwcH69iIdCWkkyjvcDwNUnTY
+    """
+    content_lst = []
+    for page_info in pdf_info_list:
+        page_lst = [] # 一个page内的段落列表
+        para_blocks = page_info.get("para_blocks")
+        pymu_raw_blocks = page_info.get("preproc_blocks")
+        
+        all_page_images = []
+        all_page_images.extend(page_info.get("images",[]))
+        all_page_images.extend(page_info.get("image_backup", []) )
+        # all_page_images.extend(page_info.get("tables",[]))
+        # all_page_images.extend(page_info.get("table_backup",[]) )
+        all_page_tables = []
+        all_page_tables.extend(page_info.get("tables", []))
+
+        if not para_blocks or not pymu_raw_blocks: # 只有图片的拼接的场景
+            for img in all_page_images:
+                content_node = {
+                    "type": "image",
+                    "img_path": join_path(img_buket_path, img['image_path']),
+                    "img_alt":"",
+                    "img_title":"",
+                    "img_caption":""
+                }
+                page_lst.append(content_node) # TODO 图片顺序
+            for table in all_page_tables:
+                content_node = {
+                    "type": "table",
+                    "img_path": join_path(img_buket_path, table['image_path']),
+                    "table_latex": table.get("text"),
+                    "table_title": "",
+                    "table_caption": "",
+                    "table_quality": table.get("quality"),
+                }
+                page_lst.append(content_node) # TODO 图片顺序
+        else:
+            for block in para_blocks:
+                item = block["paras"]
+                for _, p in item.items():
+                    font_type = p['para_font_type']# 对于文本来说，要么是普通文本，要么是个行间公式
+                    if font_type == TYPE_INTERLINE_EQUATION:
+                        content_node = {
+                            "type": "equation",
+                            "latex": p["para_text"]
+                        }
+                        page_lst.append(content_node)
+                    else:
+                        para_text = p["para_text"]
+                        is_title = p["is_para_title"]
+                        title_level = p['para_title_level']
+                        
+                        if is_title:
+                            content_node = {
+                                "type": f"h{title_level}",
+                                "text": para_text
+                            }
+                            page_lst.append(content_node)
+                        else:
+                            content_node = {
+                                "type": "text",
+                                "text": para_text
+                            }
+                            page_lst.append(content_node)
+                            
+        content_lst.extend(page_lst)
+        
+        """插入图片"""
+        for img in all_page_images:
+            insert_img_or_table("image", img, pymu_raw_blocks, content_lst)
+
+        """插入表格"""
+        for table in all_page_tables:
+            insert_img_or_table("table", table, pymu_raw_blocks, content_lst)
+    # end for
+    return content_lst
+
+
+def insert_img_or_table(type, element, pymu_raw_blocks, content_lst):
+    element_bbox = element['bbox']
+    # 先看在哪个block内
+    for block in pymu_raw_blocks:
+        bbox = block['bbox']
+        if bbox[0] - 1 <= element_bbox[0] < bbox[2] + 1 and bbox[1] - 1 <= element_bbox[1] < bbox[
+            3] + 1:  # 确定在这个大的block内，然后进入逐行比较距离
+            for l in block['lines']:
+                line_box = l['bbox']
+                if line_box[0] - 1 <= element_bbox[0] < line_box[2] + 1 and line_box[1] - 1 <= element_bbox[1] < line_box[
+                    3] + 1:  # 在line内的，插入line前面
+                    line_txt = "".join([s['text'] for s in l['spans']])
+                    __insert_before_para(line_txt, type, element, content_lst)
+                    break
+                break
+            else:  # 在行与行之间
+                # 找到图片x0,y0与line的x0,y0最近的line
+                min_distance = 100000
+                min_line = None
+                for l in block['lines']:
+                    line_box = l['bbox']
+                    distance = math.sqrt((line_box[0] - element_bbox[0]) ** 2 + (line_box[1] - element_bbox[1]) ** 2)
+                    if distance < min_distance:
+                        min_distance = distance
+                        min_line = l
+                if min_line:
+                    line_txt = "".join([s['text'] for s in min_line['spans']])
+                    img_h = element_bbox[3] - element_bbox[1]
+                    if min_distance < img_h:  # 文字在图片前面
+                        __insert_after_para(line_txt, type, element, content_lst)
+                    else:
+                        __insert_before_para(line_txt, type, element, content_lst)
+                    break
+                else:
+                    logger.error(f"Can't find the location of image {element.get('image_path')} in the markdown file #1")
+    else:  # 应当在两个block之间
+        # 找到上方最近的block，如果上方没有就找大下方最近的block
+        top_txt_block = find_top_nearest_text_bbox(pymu_raw_blocks, element_bbox)
+        if top_txt_block:
+            line_txt = "".join([s['text'] for s in top_txt_block['lines'][-1]['spans']])
+            __insert_after_para(line_txt, type, element, content_lst)
+        else:
+            bottom_txt_block = find_bottom_nearest_text_bbox(pymu_raw_blocks, element_bbox)
+            if bottom_txt_block:
+                line_txt = "".join([s['text'] for s in bottom_txt_block['lines'][0]['spans']])
+                __insert_before_para(line_txt, type, element, content_lst)
+            else:  # TODO ，图片可能独占一列，这种情况上下是没有图片的
+                logger.error(f"Can't find the location of image {element.get('image_path')} in the markdown file #2")
+
+
+def mk_mm_markdown(content_list):
+    """
+    基于同一格式的内容列表，构造markdown，含图片
+    """
+    content_md = []
+    for c in content_list:
+        content_type = c.get("type")
+        if content_type == "text":
+            content_md.append(c.get("text"))
+        elif content_type == "equation":
+            content = c.get("latex")
+            if content.startswith("$$") and content.endswith("$$"):
+                content_md.append(content)
+            else:
+                content_md.append(f"\n$$\n{c.get('latex')}\n$$\n")
+        elif content_type in UNI_FORMAT_TEXT_TYPE:
+            content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}")
+        elif content_type == "image":
+            content_md.append(f"![]({c.get('img_path')})")
+    return "\n\n".join(content_md)
+
+def mk_nlp_markdown(content_list):
+    """
+    基于同一格式的内容列表，构造markdown，不含图片
+    """
+    content_md = []
+    for c in content_list:
+        content_type = c.get("type")
+        if content_type == "text":
+            content_md.append(c.get("text"))
+        elif content_type == "equation":
+            content_md.append(f"$$\n{c.get('latex')}\n$$")
+        elif content_type == "table":
+            content_md.append(f"$$$\n{c.get('table_latex')}\n$$$")
+        elif content_type in UNI_FORMAT_TEXT_TYPE:
+            content_md.append(f"{'#'*int(content_type[1])} {c.get('text')}")
+    return "\n\n".join(content_md)
\ No newline at end of file

magic_pdf/dict2md/ocr_client.py

+import os
+import json
+import requests
+from loguru import logger
+import argparse
+import time
+from PIL import Image
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--url',
+        default='http://127.0.0.1:6020',
+    )
+    parser.add_argument(
+        '--image_path',
+        default='/home/wanglch/projects/Qwen2-VL/20240920-163701.png',
+    )
+    parser.add_argument(
+        '--text',
+        default="描述你在图片中看到的内容",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def parse_text(text):
+    lines = text.split("\n")
+    lines = [line for line in lines if line.strip() != ""]  # 去除空行
+    count = 0
+    parsed_lines = []
+
+    for i, line in enumerate(lines):
+        if "```" in line:
+            count += 1
+            items = line.split("`")
+            if count % 2 == 1:
+                # 开始代码块
+                parsed_lines.append(f'<pre><code class="language-{items[-1]}">')
+            else:
+                # 结束代码块
+                parsed_lines.append(f"</code></pre>")
+        else:
+            if i > 0 and count % 2 == 1:
+                # 转义代码块内的特殊字符
+                line = line.replace("`", r"\`")
+                line = line.replace("<", "&lt;")
+                line = line.replace(">", "&gt;")
+                line = line.replace(" ", "&nbsp;")
+                line = line.replace("*", "&ast;")
+                line = line.replace("_", "&lowbar;")
+                line = line.replace("-", "&#45;")
+                line = line.replace(".", "&#46;")
+                line = line.replace("!", "&#33;")
+                line = line.replace("(", "&#40;")
+                line = line.replace(")", "&#41;")
+                line = line.replace("$", "&#36;")
+            # 使用空格连接行
+            if parsed_lines:
+                parsed_lines[-1] += " " + line
+            else:
+                parsed_lines.append(line)
+
+    text = "".join(parsed_lines)
+    return text
+
+
+def unparse_text(parsed_text):
+    in_code_block = False
+    lines = parsed_text.split("\n")
+    unparsed_lines = []
+
+    for line in lines:
+        if "<pre><code" in line:
+            in_code_block = True
+            # 移除开始标签
+            line = line.split(">", 1)[1]
+        elif "</code></pre>" in line:
+            in_code_block = False
+            # 移除结束标签
+            line = line.rsplit("<", 1)[0]
+
+        # 反转 HTML 实体
+        line = line.replace("&lt;", "<")
+        line = line.replace("&gt;", ">")
+        line = line.replace("&nbsp;", " ")
+        line = line.replace("&ast;", "*")
+        line = line.replace("&lowbar;", "_")
+        line = line.replace("&#45;", "-")
+        line = line.replace("&#46;", ".")
+        line = line.replace("&#33;", "!")
+        line = line.replace("&#40;", "(")
+        line = line.replace("&#41;", ")")
+        line = line.replace("&#36;", "$")
+
+        # 如果在代码块内，还原反斜杠转义
+        if in_code_block:
+            line = line.replace(r"\`", "`")
+
+        unparsed_lines.append(line)
+
+    # 合并所有行
+    unparsed_text = "\n".join(unparsed_lines)
+    return unparsed_text
+
+
+def compress_image(image_path, max_size=(1024, 1024)):
+    img = Image.open(image_path)
+    width, height = img.size
+    aspect_ratio = width / height
+
+    if width > max_size[0] or height > max_size[1]:
+        if width > height:
+            new_width = max_size[0]
+            new_height = int(new_width / aspect_ratio)
+        else:
+            new_height = max_size[1]
+            new_width = int(new_height * aspect_ratio)
+
+        img = img.resize((new_width, new_height), Image.LANCZOS)
+        img.save(image_path, optimize=True, quality=80)
+
+
+class PredictClient:
+    def __init__(self, api_url):
+        self.api_url = api_url
+
+    def predict(self, image_path: str, text: str):
+        payload = {
+            "image_path": image_path,
+            "text": text
+        }
+        headers = {'Content-Type': 'application/json'}
+        response = requests.post(f"{self.api_url}/predict", json=payload, headers=headers)
+
+        if response.status_code == 200:
+            result = response.json()
+            return result.get('Generated Text', '')
+        else:
+            raise Exception(f"Predict API request failed with status code {response.status_code}")
+
+
+def main():
+    args = parse_args()
+    client = PredictClient(args.url)
+
+    try:
+        start_time = time.time()  # 记录开始时间
+        # 压缩图片
+        #compress_image(args.image_path)
+
+        generated_text = client.predict(args.image_path, parse_text(args.text))
+        end_time = time.time()  # 记录结束时间
+        elapsed_time = end_time - start_time  # 计算运行时间
+
+        if generated_text:
+            clean_text = unparse_text(generated_text)  # 解析生成的文本
+            logger.info(f"Image Path: {args.image_path}")
+            logger.info(f"Generated Text: {clean_text}")
+            logger.info(f"耗时为: {elapsed_time}秒")  # 打印运行时间
+        else:
+            logger.warning("Received empty generated text.")
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Error while making request to predict service: {e}")
+    except Exception as e:
+        logger.error(f"Unexpected error occurred: {e}")
+
+
+if __name__ == "__main__":
+    main()

magic_pdf/dict2md/ocr_server.py

+# Copyright (c) Alibaba Cloud.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree
+
+import copy
+import re
+import gc
+import torch
+from argparse import ArgumentParser
+from threading import Thread
+from qwen_vl_utils import process_vision_info
+from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, TextIteratorStreamer
+
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import Optional
+
+app = FastAPI()
+
+DEFAULT_CKPT_PATH = '/home/practice/model/Qwen2-VL-7B-Instruct'
+REVISION = 'v1.0.4'
+BOX_TAG_PATTERN = r"<box>([\s\S]*?)</box>"
+PUNCTUATION = "！？。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."
+
+
+def _get_args():
+    parser = ArgumentParser()
+
+    parser.add_argument('-c', '--checkpoint_path', type=str, default=DEFAULT_CKPT_PATH,
+                        help='Checkpoint name or path, default to %(default)r')
+    parser.add_argument('--cpu_only', action='store_true', help='Run demo with CPU only')
+    parser.add_argument('--flash_attn2', action='store_true', default=False,
+                        help='Enable flash_attention_2 when loading the model.')
+    parser.add_argument('--share', action='store_true', default=False,
+                        help='Create a publicly shareable link for the interface.')
+    parser.add_argument('--inbrowser', action='store_true', default=False,
+                        help='Automatically launch the interface in a new tab on the default browser.')
+    parser.add_argument('--server_port', type=int, default=6020, help='Demo server port.')
+    parser.add_argument('--server_name', type=str, default='127.0.0.1', help='Demo server name.')
+    parser.add_argument('--dcu_id', type=str, default='0', help='Specify the GPU ID to load the model onto.')
+
+    args = parser.parse_args()
+    return args
+
+
+def _load_model_processor(args):
+    if args.cpu_only:
+        device_map = 'cpu'
+    else:
+        if args.dcu_id is not None:
+            device_map = {'': f'cuda:{args.dcu_id}'}
+            print('使用DCU推理:', f'cuda:{args.dcu_id}')
+        else:
+            device_map = 'auto'
+
+    if args.flash_attn2:
+        model = Qwen2VLForConditionalGeneration.from_pretrained(
+            args.checkpoint_path,
+            torch_dtype=torch.float16,
+            attn_implementation='flash_attention_2',
+            device_map=device_map
+        )
+    else:
+        model = Qwen2VLForConditionalGeneration.from_pretrained(
+            args.checkpoint_path,
+            torch_dtype=torch.float16,
+            device_map=device_map
+        )
+
+    processor = AutoProcessor.from_pretrained(args.checkpoint_path)
+    return model, processor
+
+
+def _parse_text(text):
+    lines = text.split("\n")
+    lines = [line for line in lines if line.strip() != ""]  # 去除空行
+    count = 0
+    parsed_lines = []
+
+    for i, line in enumerate(lines):
+        if "```" in line:
+            count += 1
+            items = line.split("`")
+            if count % 2 == 1:
+                # 开始代码块
+                parsed_lines.append(f'<pre><code class="language-{items[-1]}">')
+            else:
+                # 结束代码块
+                parsed_lines.append(f"</code></pre>")
+        else:
+            if i > 0 and count % 2 == 1:
+                # 转义代码块内的特殊字符
+                line = line.replace("`", r"\`")
+                line = line.replace("<", "&lt;")
+                line = line.replace(">", "&gt;")
+                line = line.replace(" ", "&nbsp;")
+                line = line.replace("*", "&ast;")
+                line = line.replace("_", "&lowbar;")
+                line = line.replace("-", "&#45;")
+                line = line.replace(".", "&#46;")
+                line = line.replace("!", "&#33;")
+                line = line.replace("(", "&#40;")
+                line = line.replace(")", "&#41;")
+                line = line.replace("$", "&#36;")
+            # 使用空格连接行
+            if parsed_lines:
+                parsed_lines[-1] += " " + line
+            else:
+                parsed_lines.append(line)
+
+    text = "".join(parsed_lines)
+    return text
+
+
+def _remove_image_special(text):
+    text = text.replace('<ref>', '').replace('</ref>', '')
+    return re.sub(r'<box>.*?(</box>|$)', '', text)
+
+
+def _is_video_file(filename):
+    video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.mpeg']
+    return any(filename.lower().endswith(ext) for ext in video_extensions)
+
+
+def _transform_messages(original_messages):
+    transformed_messages = []
+    for message in original_messages:
+        new_content = []
+        for item in message['content']:
+            if 'image' in item:
+                new_item = {'type': 'image', 'image': item['image']}
+            elif 'text' in item:
+                new_item = {'type': 'text', 'text': item['text']}
+            elif 'video' in item:
+                new_item = {'type': 'video', 'video': item['video']}
+            else:
+                continue
+            new_content.append(new_item)
+
+        new_message = {'role': message['role'], 'content': new_content}
+        transformed_messages.append(new_message)
+
+    return transformed_messages
+
+
+def _gc():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def call_local_model(model, processor, messages):
+    messages = _transform_messages(messages)
+
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors='pt')
+    inputs = inputs.to(model.device)
+
+    tokenizer = processor.tokenizer
+    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
+
+    gen_kwargs = {'max_new_tokens': 512, 'streamer': streamer, **inputs}
+
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
+
+    generated_text = ''
+    for new_text in streamer:
+        generated_text += new_text
+        yield _parse_text(generated_text)
+
+
+def create_predict_fn(model, processor):
+    def predict(_chatbot, task_history):
+        chat_query = _chatbot[-1][0]
+        query = task_history[-1][0]
+        if len(chat_query) == 0:
+            _chatbot.pop()
+            task_history.pop()
+            return _chatbot
+        print('User: ' + _parse_text(query))
+        history_cp = copy.deepcopy(task_history)
+        full_response = ''
+        messages = []
+        content = []
+        for q, a in history_cp:
+            if isinstance(q, (tuple, list)):
+                if _is_video_file(q[0]):
+                    content.append({'video': f'file://{q[0]}'})
+                else:
+                    content.append({'image': f'file://{q[0]}'})
+            else:
+                content.append({'text': q})
+                messages.append({'role': 'user', 'content': content})
+                messages.append({'role': 'assistant', 'content': [{'text': a}]})
+                content = []
+        messages.pop()
+
+        for response in call_local_model(model, processor, messages):
+            _chatbot[-1] = (_parse_text(chat_query), _remove_image_special(_parse_text(response)))
+
+            yield _chatbot
+            full_response = _parse_text(response)
+
+        task_history[-1] = (query, full_response)
+        print('Qwen-VL-Chat: ' + _parse_text(full_response))
+        yield _chatbot
+
+    return predict
+
+
+# 启用加载模型
+args = _get_args()
+model, processor = _load_model_processor(args)
+
+
+class Item(BaseModel):
+    image_path: str
+    text: str
+
+
+@app.post("/predict")
+async def predict(item: Item):
+    messages = [
+        {
+            'role': 'user',
+            'content': [
+                {'image': item.image_path},
+                {'text': item.text}
+            ]
+        }
+    ]
+
+    generated_text = ''
+    for response in call_local_model(model, processor, messages):
+        generated_text = _parse_text(response)
+
+    _gc()
+
+    return {"Generated Text": generated_text}
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    args = _get_args()
+    uvicorn.run(app, host=args.server_name, port=args.server_port)
+
+

magic_pdf/dict2md/ocr_server_72.py

+# Copyright (c) Alibaba Cloud.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree
+
+import copy
+import re
+import gc
+import torch
+from argparse import ArgumentParser
+from threading import Thread
+from qwen_vl_utils import process_vision_info
+from transformers import AutoProcessor, Qwen2VLForConditionalGeneration, TextIteratorStreamer
+import os
+
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import Optional
+
+app = FastAPI()
+
+DEFAULT_CKPT_PATH = '/home/practice/model/Qwen2-VL-72B-Instruct'
+REVISION = 'v1.0.4'
+BOX_TAG_PATTERN = r"<box>([\s\S]*?)</box>"
+PUNCTUATION = "！？。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏."
+
+def get_args():
+    parser = ArgumentParser()
+
+    parser.add_argument('-c', '--checkpoint_path', type=str, default=DEFAULT_CKPT_PATH,
+                        help='Checkpoint name or path, default to %(default)r')
+    parser.add_argument('--cpu_only', action='store_true', help='Run demo with CPU only')
+    parser.add_argument('--flash_attn2', action='store_true', default=False,
+                        help='Enable flash_attention_2 when loading the model.')
+    parser.add_argument('--share', action='store_true', default=False,
+                        help='Create a publicly shareable link for the interface.')
+    parser.add_argument('--inbrowser', action='store_true', default=False,
+                        help='Automatically launch the interface in a new tab on the default browser.')
+    parser.add_argument('--server_port', type=int, default=6020, help='Demo server port.')
+    parser.add_argument('--server_name', type=str, default='127.0.0.1', help='Demo server name.')
+    parser.add_argument('--dcu_id', type=str, default=None, help='Specify the GPU ID to load the model onto.')
+
+    args = parser.parse_args()
+    return args
+
+
+def load_model_processor(args):
+    if args.cpu_only:
+        device_map = 'cpu'
+    else:
+        os.environ['CUDA_VISIBLE_DEVICES'] = args.dcu_id
+        device_map = 'auto'
+        print(f"Visible CUDA devices: {os.environ.get('CUDA_VISIBLE_DEVICES')}")
+
+    if args.flash_attn2:
+        model = Qwen2VLForConditionalGeneration.from_pretrained(
+            args.checkpoint_path,
+            torch_dtype=torch.float16,
+            attn_implementation='flash_attention_2',
+            device_map=device_map
+        )
+    else:
+        model = Qwen2VLForConditionalGeneration.from_pretrained(
+            args.checkpoint_path,
+            torch_dtype=torch.float16,
+            device_map=device_map
+        )
+
+    processor = AutoProcessor.from_pretrained(args.checkpoint_path)
+    return model, processor
+
+def parse_text(text):
+    lines = text.split("\n")
+    lines = [line for line in lines if line.strip() != ""]  # 去除空行
+    count = 0
+    parsed_lines = []
+
+    for i, line in enumerate(lines):
+        if "```" in line:
+            count += 1
+            items = line.split("`")
+            if count % 2 == 1:
+                # 开始代码块
+                parsed_lines.append(f'<pre><code class="language-{items[-1]}">')
+            else:
+                # 结束代码块
+                parsed_lines.append(f"</code></pre>")
+        else:
+            if i > 0 and count % 2 == 1:
+                # 转义代码块内的特殊字符
+                line = line.replace("`", r"\`")
+                line = line.replace("<", "&lt;")
+                line = line.replace(">", "&gt;")
+                line = line.replace(" ", "&nbsp;")
+                line = line.replace("*", "&ast;")
+                line = line.replace("_", "&lowbar;")
+                line = line.replace("-", "&#45;")
+                line = line.replace(".", "&#46;")
+                line = line.replace("!", "&#33;")
+                line = line.replace("(", "&#40;")
+                line = line.replace(")", "&#41;")
+                line = line.replace("$", "&#36;")
+            # 使用空格连接行
+            if parsed_lines:
+                parsed_lines[-1] += " " + line
+            else:
+                parsed_lines.append(line)
+
+    text = "".join(parsed_lines)
+    return text
+
+def unparse_text(parsed_text):
+    in_code_block = False
+    lines = parsed_text.split("\n")
+    unparsed_lines = []
+    
+    for line in lines:
+        if "<pre><code" in line:
+            in_code_block = True
+            # 移除开始标签
+            line = line.split(">", 1)[1]
+        elif "</code></pre>" in line:
+            in_code_block = False
+            # 移除结束标签
+            line = line.rsplit("<", 1)[0]
+        
+        # 反转 HTML 实体
+        line = line.replace("&lt;", "<")
+        line = line.replace("&gt;", ">")
+        line = line.replace("&nbsp;", " ")
+        line = line.replace("&ast;", "*")
+        line = line.replace("&lowbar;", "_")
+        line = line.replace("&#45;", "-")
+        line = line.replace("&#46;", ".")
+        line = line.replace("&#33;", "!")
+        line = line.replace("&#40;", "(")
+        line = line.replace("&#41;", ")")
+        line = line.replace("&#36;", "$")
+        
+        # 如果在代码块内，还原反斜杠转义
+        if in_code_block:
+            line = line.replace(r"\`", "`")
+        
+        unparsed_lines.append(line)
+    
+    # 合并所有行
+    unparsed_text = "\n".join(unparsed_lines)
+    return unparsed_text
+
+def remove_image_special(text):
+    text = text.replace('<ref>', '').replace('</ref>', '')
+    return re.sub(r'<box>.*?(</box>|$)', '', text)
+
+def is_video_file(filename):
+    video_extensions = ['.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.mpeg']
+    return any(filename.lower().endswith(ext) for ext in video_extensions)
+
+def transform_messages(original_messages):
+    transformed_messages = []
+    for message in original_messages:
+        new_content = []
+        for item in message['content']:
+            if 'image' in item:
+                new_item = {'type': 'image', 'image': item['image']}
+            elif 'text' in item:
+                new_item = {'type': 'text', 'text': item['text']}
+            elif 'video' in item:
+                new_item = {'type': 'video', 'video': item['video']}
+            else:
+                continue
+            new_content.append(new_item)
+
+        new_message = {'role': message['role'], 'content': new_content}
+        transformed_messages.append(new_message)
+
+    return transformed_messages
+
+def _gc():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+def call_local_model(model, processor, messages):
+    messages = transform_messages(messages)
+
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors='pt')
+    inputs = inputs.to(model.device)
+
+    tokenizer = processor.tokenizer
+    streamer = TextIteratorStreamer(tokenizer, timeout=2000.0, skip_prompt=True, skip_special_tokens=True)
+
+    gen_kwargs = {'max_new_tokens': 512, 'streamer': streamer, **inputs}
+
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
+
+    generated_text = ''
+    for new_text in streamer:
+        generated_text += new_text
+        yield parse_text(generated_text)
+
+def create_predict_fn(model, processor):
+    def predict(_chatbot, task_history):
+        chat_query = _chatbot[-1][0]
+        query = task_history[-1][0]
+        if len(chat_query) == 0:
+            _chatbot.pop()
+            task_history.pop()
+            return _chatbot
+        print('User: ' + parse_text(query))
+        history_cp = copy.deepcopy(task_history)
+        full_response = ''
+        messages = []
+        content = []
+        for q, a in history_cp:
+            if isinstance(q, (tuple, list)):
+                if is_video_file(q[0]):
+                    content.append({'video': f'file://{q[0]}'})
+                else:
+                    content.append({'image': f'file://{q[0]}'})
+            else:
+                content.append({'text': q})
+                messages.append({'role': 'user', 'content': content})
+                messages.append({'role': 'assistant', 'content': [{'text': a}]})
+                content = []
+        messages.pop()
+
+        for response in call_local_model(model, processor, messages):
+            _chatbot[-1] = (parse_text(chat_query), remove_image_special(parse_text(response)))
+
+            yield _chatbot
+            full_response = parse_text(response)
+
+        task_history[-1] = (query, full_response)
+        print('Qwen-VL-Chat: ' + unparse_text(full_response))
+        yield _chatbot
+
+    return predict
+
+# 启用加载模型
+args = get_args()
+model, processor = load_model_processor(args)
+
+class Item(BaseModel):
+    image_path: str
+    text: str
+
+@app.post("/predict")
+async def predict(item: Item):
+    messages = [
+        {
+            'role': 'user',
+            'content': [
+                {'image': item.image_path},
+                {'text': item.text}
+            ]
+        }
+    ]
+
+    generated_text = ''
+    for response in call_local_model(model, processor, messages):
+        generated_text = unparse_text(response)
+
+    _gc()
+
+    return {"Generated Text": generated_text}
+
+if __name__ == "__main__":
+    import uvicorn
+    args = get_args()
+    uvicorn.run(app, host=args.server_name, port=args.server_port)

magic_pdf/dict2md/tmp.py

+# -*- coding: utf-8 -*-
+import time
+import os
+
+import ocr_client
+text = '识别图片的内容，如果是发票就执行以下操作识别图中的文字信息，并以json格式返回，如果不是发票返回False'
+image_dir = '/home/practice/dcu_Mineru/images'
+false_count = 0
+file_count = 0
+true_count = 0
+for root, dirs, files in os.walk(image_dir):
+    for file in files:
+        # 构建完整的文件路径
+        file_path = os.path.join(root, file)
+        # 打印文件名
+        image_path = file_path
+        start = time.time()
+        ocr_client.compress_image(image_path)
+        generated_text = ocr_client.PredictClient('http://127.0.0.1:6020').predict(image_path, text)
+
+        if generated_text == 'False' or generated_text == 'false':
+            false_count += 1
+        else:
+            print(f'{image_path}提取到的内容是：\n{generated_text}')
+            true_count += 1
+        file_count +=1
+
+print(f'总共{file_count}个文件，其中qwen返回{false_count}个False，返回{true_count}个json')
+
+
+