Merge pull request #2625 from opendatalab/release-2.0.0

Release 2.0.0

magic_pdf/libs/clean_memory.py

-# Copyright (c) Opendatalab. All rights reserved.
-import torch
-import gc
-
-
-def clean_memory(device='cuda'):
-    if device == 'cuda':
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            torch.cuda.ipc_collect()
-    elif str(device).startswith("npu"):
-        import torch_npu
-        if torch_npu.npu.is_available():
-            torch_npu.npu.empty_cache()
-    elif str(device).startswith("mps"):
-        torch.mps.empty_cache()
-    gc.collect()
\ No newline at end of file

magic_pdf/libs/commons.py

-
-def join_path(*args):
-    return '/'.join(str(s).rstrip('/') for s in args)
-
-
-def get_top_percent_list(num_list, percent):
-    """
-    获取列表中前百分之多少的元素
-    :param num_list:
-    :param percent:
-    :return:
-    """
-    if len(num_list) == 0:
-        top_percent_list = []
-    else:
-        # 对imgs_len_list排序
-        sorted_imgs_len_list = sorted(num_list, reverse=True)
-        # 计算 percent 的索引
-        top_percent_index = int(len(sorted_imgs_len_list) * percent)
-        # 取前80%的元素
-        top_percent_list = sorted_imgs_len_list[:top_percent_index]
-    return top_percent_list
-
-
-def mymax(alist: list):
-    if len(alist) == 0:
-        return 0  # 空是0， 0*0也是0大小q
-    else:
-        return max(alist)
-
-
-def parse_bucket_key(s3_full_path: str):
-    """
-    输入 s3://bucket/path/to/my/file.txt
-    输出 bucket, path/to/my/file.txt
-    """
-    s3_full_path = s3_full_path.strip()
-    if s3_full_path.startswith("s3://"):
-        s3_full_path = s3_full_path[5:]
-    if s3_full_path.startswith("/"):
-        s3_full_path = s3_full_path[1:]
-    bucket, key = s3_full_path.split("/", 1)
-    return bucket, key

magic_pdf/libs/convert_utils.py

-def dict_to_list(input_dict):
-    items_list = []
-    for _, item in input_dict.items():
-        items_list.append(item)
-    return items_list

magic_pdf/libs/coordinate_transform.py

-def get_scale_ratio(model_page_info, page):
-    pix = page.get_pixmap(dpi=72)
-    pymu_width = int(pix.w)
-    pymu_height = int(pix.h)
-    width_from_json = model_page_info['page_info']['width']
-    height_from_json = model_page_info['page_info']['height']
-    horizontal_scale_ratio = width_from_json / pymu_width
-    vertical_scale_ratio = height_from_json / pymu_height
-    return horizontal_scale_ratio, vertical_scale_ratio

magic_pdf/libs/json_compressor.py

-import json
-import brotli
-import base64
-
-class JsonCompressor:
-
-    @staticmethod
-    def compress_json(data):
-        """
-        Compress a json object and encode it with base64
-        """
-        json_str = json.dumps(data)
-        json_bytes = json_str.encode('utf-8')
-        compressed = brotli.compress(json_bytes, quality=6)
-        compressed_str = base64.b64encode(compressed).decode('utf-8')  # convert bytes to string
-        return compressed_str
-
-    @staticmethod
-    def decompress_json(compressed_str):
-        """
-        Decode the base64 string and decompress the json object
-        """
-        compressed = base64.b64decode(compressed_str.encode('utf-8'))  # convert string to bytes
-        decompressed_bytes = brotli.decompress(compressed)
-        json_str = decompressed_bytes.decode('utf-8')
-        data = json.loads(json_str)
-        return data

magic_pdf/libs/local_math.py

-def float_gt(a, b):
-    if 0.0001 >= abs(a -b):
-        return False
-    return a > b
-    
-def float_equal(a, b):
-    if 0.0001 >= abs(a-b):
-        return True
-    return False
\ No newline at end of file

magic_pdf/libs/markdown_utils.py

-
-def ocr_escape_special_markdown_char(content):
-    """
-    转义正文里对markdown语法有特殊意义的字符
-    """
-    special_chars = ["*", "`", "~", "$"]
-    for char in special_chars:
-        content = content.replace(char, "\\" + char)
-
-    return content

magic_pdf/libs/pdf_image_tools.py

-from io import BytesIO
-import cv2
-import fitz
-import numpy as np
-from PIL import Image
-from magic_pdf.data.data_reader_writer import DataWriter
-from magic_pdf.libs.commons import join_path
-from magic_pdf.libs.hash_utils import compute_sha256
-
-
-def cut_image(bbox: tuple, page_num: int, page: fitz.Page, return_path, imageWriter: DataWriter):
-    """从第page_num页的page中，根据bbox进行裁剪出一张jpg图片，返回图片路径 save_path：需要同时支持s3和本地,
-    图片存放在save_path下，文件名是:
-    {page_num}_{bbox[0]}_{bbox[1]}_{bbox[2]}_{bbox[3]}.jpg , bbox内数字取整。"""
-    # 拼接文件名
-    filename = f'{page_num}_{int(bbox[0])}_{int(bbox[1])}_{int(bbox[2])}_{int(bbox[3])}'
-
-    # 老版本返回不带bucket的路径
-    img_path = join_path(return_path, filename) if return_path is not None else None
-
-    # 新版本生成平铺路径
-    img_hash256_path = f'{compute_sha256(img_path)}.jpg'
-
-    # 将坐标转换为fitz.Rect对象
-    rect = fitz.Rect(*bbox)
-    # 配置缩放倍数为3倍
-    zoom = fitz.Matrix(3, 3)
-    # 截取图片
-    pix = page.get_pixmap(clip=rect, matrix=zoom)
-
-    byte_data = pix.tobytes(output='jpeg', jpg_quality=95)
-
-    imageWriter.write(img_hash256_path, byte_data)
-
-    return img_hash256_path
-
-
-def cut_image_to_pil_image(bbox: tuple, page: fitz.Page, mode="pillow"):
-
-    # 将坐标转换为fitz.Rect对象
-    rect = fitz.Rect(*bbox)
-    # 配置缩放倍数为3倍
-    zoom = fitz.Matrix(3, 3)
-    # 截取图片
-    pix = page.get_pixmap(clip=rect, matrix=zoom)
-
-    if mode == "cv2":
-        # 直接转换为numpy数组供cv2使用
-        img_array = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
-        # PyMuPDF使用RGB顺序，而cv2使用BGR顺序
-        if pix.n == 3 or pix.n == 4:
-            image_result = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
-        else:
-            image_result = img_array
-    elif mode == "pillow":
-        # 将字节数据转换为文件对象
-        image_file = BytesIO(pix.tobytes(output='png'))
-        # 使用 Pillow 打开图像
-        image_result = Image.open(image_file)
-    else:
-        raise ValueError(f"mode: {mode} is not supported.")
-
-    return image_result
\ No newline at end of file

magic_pdf/libs/version.py

-__version__ = "1.3.12"

magic_pdf/model/__init__.py

-__use_inside_model__ = True
-__model_mode__ = 'full'
\ No newline at end of file