pydatamax 0.1.5__py3-none-any.whl → 0.1.12__py3-none-any.whl

datamax/utils/env_setup.py

@@ -0,0 +1,80 @@
+import subprocess
+import sys
+import os
+import importlib.metadata
+
+class EnvironmentSetup:
+    """ Responsible for setting up the correct environment,
+        including checking GPU support and installing the necessary packages
+    """
+
+    def __init__(self, use_gpu: bool = False):
+        self._gpu_available = None
+        self._setup_completed = False
+        self.use_gpu = use_gpu  # Use GPU if True, otherwise use CPU
+
+    def is_gpu_available(self):
+        """Check whether the system supports Gpus"""
+        if self._gpu_available is None:
+            try:
+                # Check whether CUDA is available
+                subprocess.check_output(['nvcc', '--version'], stderr=subprocess.STDOUT)
+                self._gpu_available = True
+            except (subprocess.CalledProcessError, FileNotFoundError):
+                self._gpu_available = False
+        return self._gpu_available
+
+    def is_conda(self):
+        """ Check whether the current environment is a Conda environment """
+        return os.path.exists(os.path.join(sys.prefix, 'conda-meta'))
+
+    def install_package(self, package_name):
+        """ Select pip or conda or other installation specified package according to the environment """
+        installer = 'conda' if self.is_conda() else 'pip'
+        if installer == 'conda':
+            print(f"Detected Conda environment. Installing {package_name} with conda.")
+            try:
+                subprocess.check_call(['pip', 'install', package_name])
+                print(f"Successfully installed {package_name} with conda.")
+            except subprocess.CalledProcessError as e:
+                print(f"Failed to install {package_name} with conda: {e}")
+        elif installer == 'pip':
+            print(f"Using pip to install {package_name}.")
+            try:
+                # Invoke the pip installation package using the Python interpreter
+                subprocess.check_call([sys.executable, "-m", "pip", "install", package_name])
+                print(f"Successfully installed {package_name} with pip.")
+            except subprocess.CalledProcessError as e:
+                print(f"Failed to install {package_name} with pip: {e}")
+        else:
+            print("Unable to determine the package manager. Please install the package manually.")
+
+    def check_and_install(self):
+        """Check and install appropriate packages based on user's choice and GPU availability"""
+        if self._setup_completed:
+            return
+
+        # Override GPU detection with the use_gpu parameter
+        if self.use_gpu:
+            pkg_name = 'paddlepaddle-gpu' if self.is_gpu_available() else 'paddlepaddle'
+        else:
+            pkg_name = 'paddlepaddle'
+
+        try:
+            _ = importlib.metadata.version(pkg_name.split()[0])  # Check if paddlepaddle is installed
+            # print(f"{pkg_name} version {1} is already installed.")
+        except importlib.metadata.PackageNotFoundError:
+            print(f"{pkg_name} is not installed. Installing now...")
+            self.install_package(pkg_name)
+
+        self._setup_completed = True
+
+
+# Create an instance of EnvironmentSetup with the desired GPU option and call check_and_install when the program initializes
+env_setup = EnvironmentSetup()  # Set this flag as needed
+
+
+def setup_environment(use_gpu: bool = False):
+    """Used to set the environment when the program starts"""
+    env_setup.use_gpu = use_gpu
+    env_setup.check_and_install()

datamax/utils/gotocr_pdf.py

@@ -0,0 +1,265 @@
+import argparse
+import ast
+import os
+import re
+from datetime import datetime
+
+import cv2
+import numpy as np
+import torch
+from GOT.demo.process_results import punctuation_dict
+from GOT.model import *
+from GOT.model.plug.blip_process import BlipImageEvalProcessor
+from GOT.utils.conversation import SeparatorStyle, conv_templates
+from GOT.utils.utils import KeywordsStoppingCriteria, disable_torch_init
+from paddle.utils import try_import
+from PIL import Image
+from transformers import AutoTokenizer
+
+fitz = try_import("fitz")
+
+DEFAULT_IMAGE_TOKEN = "<image>"  # nosec B105 - 这是技术常量，不是密码
+DEFAULT_IMAGE_PATCH_TOKEN = "<imgpad>"  # nosec B105 - 这是技术常量，不是密码
+
+DEFAULT_IM_START_TOKEN = "<img>"  # nosec B105 - 这是技术常量，不是密码
+DEFAULT_IM_END_TOKEN = "</img>"  # nosec B105 - 这是技术常量，不是密码
+
+translation_table = str.maketrans(punctuation_dict)
+
+parser = argparse.ArgumentParser()
+
+args = argparse.Namespace()
+args.model_name = "./GOT_weights/"
+args.type = "format"
+args.box = ""
+args.color = ""
+
+# TODO vary old codes, NEED del
+image_processor = BlipImageEvalProcessor(image_size=1024)
+
+image_processor_high = BlipImageEvalProcessor(image_size=1024)
+
+use_im_start_end = True
+
+image_token_len = 256
+
+
+def covert_pdf_to_image(image_path: str):
+    # step1: Convert PDF to images
+    imgs = []
+    with fitz.open(image_path) as pdf:
+        for pg in range(0, pdf.page_count):
+            page = pdf[pg]
+            mat = fitz.Matrix(4, 4)  # 全程放大四倍
+            pm = page.get_pixmap(matrix=mat, alpha=False)
+            # if pm.width > 2000 or pm.height > 2000:
+            #     pm = page.get_pixmap(matrix=fitz.Matrix(1, 1), alpha=False)
+
+            img = Image.frombytes("RGB", [pm.width, pm.height], pm.samples)
+            img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+            imgs.append(img)
+
+    img_name = datetime.now().strftime("%Y%m%d%H%M%S")
+    # step2: Process images
+    output = "output"
+    img_paths = []
+    for index, pdf_img in enumerate(imgs):
+        # 图片处理
+
+        gray_img = cv2.cvtColor(pdf_img, cv2.COLOR_BGR2GRAY)
+
+        # 二值化处理
+        _, binary_img = cv2.threshold(
+            gray_img, 128, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU
+        )
+
+        # 去噪
+        filtered_img = cv2.medianBlur(binary_img, 3)
+        processed_img = filtered_img
+
+        os.makedirs(os.path.join(output, img_name), exist_ok=True)
+        pdf_img_path = os.path.join(
+            output, img_name, img_name + "_" + str(index) + ".jpg"
+        )
+        cv2.imwrite(pdf_img_path, processed_img)
+        img_paths.append([pdf_img_path, processed_img])
+
+    return img_name
+
+
+# def initialize_model(model_path: str = "./GOT_weights/", gpu_id: int = 6):
+#     model_name = os.path.expanduser(model_path)
+#     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+#     model = GOTQwenForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True, device_map=f'cuda:{gpu_id}',
+#                                                use_safetensors=True, pad_token_id=151643).eval()
+#     model.to(device=f'cuda:{gpu_id}', dtype=torch.bfloat16)
+#     return model, tokenizer
+
+
+def initialize_model(model_path: str = "./GOT_weights/", gpu_id: int = 6):
+    model_name = os.path.expanduser(model_path)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+    # 加载模型
+    model = GOTQwenForCausalLM.from_pretrained(
+        model_name,
+        low_cpu_mem_usage=True,
+        device_map=f"cuda:{gpu_id}",
+        use_safetensors=True,
+        pad_token_id=151643,
+    ).eval()
+
+    # 确保模型和张量都移动到目标设备
+    device = torch.device(f"cuda:{gpu_id}")
+    model.to(device=device, dtype=torch.bfloat16)
+
+    # 确保分词器的输出也在目标设备上
+    tokenizer.model_max_length = 512  # 设置最大长度，根据需要调整
+    tokenizer.padding_side = "right"  # 设置填充方向，根据需要调整
+
+    return model, tokenizer
+
+
+def eval_model(file: str, model, tokenizer, gpu_id: int = 6):
+    # Model
+    # image = load_image(args.image_file)
+    image = Image.open(file).convert("RGB")
+
+    w, h = image.size
+    # print(image.size)
+
+    disable_torch_init()
+
+    if args.type == "format":
+        qs = "OCR with format: "
+    else:
+        qs = "OCR: "
+
+    if args.box:
+        bbox = ast.literal_eval(args.box)
+        if len(bbox) == 2:
+            bbox[0] = int(bbox[0] / w * 1000)
+            bbox[1] = int(bbox[1] / h * 1000)
+        if len(bbox) == 4:
+            bbox[0] = int(bbox[0] / w * 1000)
+            bbox[1] = int(bbox[1] / h * 1000)
+            bbox[2] = int(bbox[2] / w * 1000)
+            bbox[3] = int(bbox[3] / h * 1000)
+        if args.type == "format":
+            qs = str(bbox) + " " + "OCR with format: "
+        else:
+            qs = str(bbox) + " " + "OCR: "
+
+    if args.color:
+        if args.type == "format":
+            qs = "[" + args.color + "]" + " " + "OCR with format: "
+        else:
+            qs = "[" + args.color + "]" + " " + "OCR: "
+
+    if use_im_start_end:
+        qs = (
+            DEFAULT_IM_START_TOKEN
+            + DEFAULT_IMAGE_PATCH_TOKEN * image_token_len
+            + DEFAULT_IM_END_TOKEN
+            + "\n"
+            + qs
+        )
+    else:
+        qs = DEFAULT_IMAGE_TOKEN + "\n" + qs
+
+    conv_mode = "mpt"
+    args.conv_mode = conv_mode
+
+    conv = conv_templates[args.conv_mode].copy()
+    conv.append_message(conv.roles[0], qs)
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt()
+
+    inputs = tokenizer([prompt])
+
+    # vary old codes, no use
+    image_1 = image.copy()
+    image_tensor = image_processor(image)
+
+    image_tensor_1 = image_processor_high(image_1)
+
+    input_ids = torch.as_tensor(inputs.input_ids).to(f"cuda:{gpu_id}")
+
+    stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+    keywords = [stop_str]
+    stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+    # streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    with torch.autocast(f"cuda:{gpu_id}", dtype=torch.bfloat16):
+        output_ids = model.generate(
+            input_ids,
+            images=[
+                (
+                    image_tensor.unsqueeze(0).half().to(f"cuda:{gpu_id}"),
+                    image_tensor_1.unsqueeze(0).half().to(f"cuda:{gpu_id}"),
+                )
+            ],
+            do_sample=False,
+            num_beams=1,
+            no_repeat_ngram_size=20,
+            # streamer=streamer,
+            max_new_tokens=4096,
+            stopping_criteria=[stopping_criteria],
+        )
+
+        outputs = tokenizer.decode(output_ids[0, input_ids.shape[1] :]).strip()
+
+        if outputs.endswith(stop_str):
+            outputs = outputs[: -len(stop_str)]
+        outputs = outputs.strip()
+    return outputs + "\n"
+
+
+def sorted_list_by_index(img_name):
+    files_with_index = []
+    for root, dirs, files in os.walk(f"./output/{img_name}"):
+        for file in files:
+            file_path = os.path.join(root, file)
+            match = re.search(r"_(\d+)(?:\.\w+)?$", file)
+            if match:
+                index = int(match.group(1))
+                files_with_index.append((file_path, index))
+    files_with_index.sort(key=lambda x: x[1])
+    sorted_files = [file[0] for file in files_with_index]
+    return sorted_files
+
+
+def convert_to_markdown(md_content, pdf_path):
+    """write into markdown"""
+    file_extension = os.path.splitext(pdf_path)[1].lower()
+    output_file = (
+        f'./got_output/{os.path.basename(pdf_path).replace(file_extension, ".mmd")}'
+    )
+    os.makedirs("got_output", exist_ok=True)
+    with open(output_file, "w", encoding="utf-8") as f:
+        f.write(md_content)
+
+
+def main(image_list: str, pdf_path: str, model, tokenizer, gpu_id: int = 6):
+    res_list = sorted_list_by_index(image_list)
+
+    outputs = ""
+    for file_path in res_list:
+        outputs += eval_model(
+            file=file_path, model=model, tokenizer=tokenizer, gpu_id=gpu_id
+        )
+
+    convert_to_markdown(outputs, pdf_path)
+    return outputs
+
+
+def generate_mathpix_markdown(pdf_path: str, model, tokenizer, gpu_id: int = 6):
+    image_list = covert_pdf_to_image(pdf_path)
+    outputs = main(
+        image_list=image_list,
+        pdf_path=pdf_path,
+        gpu_id=gpu_id,
+        model=model,
+        tokenizer=tokenizer,
+    )
+    return outputs

datamax/utils/mineru_operator.py

@@ -0,0 +1,62 @@
+import os
+from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
+from magic_pdf.data.dataset import PymuDocDataset
+from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
+from magic_pdf.config.enums import SupportedPdfParseMethod
+
+
+class PdfProcessor:
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(PdfProcessor, cls).__new__(cls)
+        return cls._instance
+
+    def process_pdf(self, pdf_file_name):
+        name_without_suff = os.path.basename(pdf_file_name).split(".")[0]
+        print("Processing PDF: " + name_without_suff)
+
+        local_image_dir = "uploaded_files/images"
+        local_md_dir = "uploaded_files/markdown"
+
+        os.makedirs(local_image_dir, exist_ok=True)
+        os.makedirs(local_md_dir, exist_ok=True)
+
+        image_writer = FileBasedDataWriter(local_image_dir)
+        md_writer = FileBasedDataWriter(local_md_dir)
+
+        reader = FileBasedDataReader("")
+        pdf_bytes = reader.read(pdf_file_name)
+
+        # 处理流程
+        ds = PymuDocDataset(pdf_bytes)
+        markdown_path = os.path.join(local_md_dir, f"{name_without_suff}.md")  # 完整路径
+        image_dir = os.path.basename(local_image_dir)  # 保持相对路径为 "images"
+
+        if ds.classify() == SupportedPdfParseMethod.OCR:
+            ds.apply(doc_analyze, ocr=True).pipe_ocr_mode(image_writer).dump_md(
+                md_writer,
+                os.path.basename(markdown_path),  # 文件名部分
+                image_dir
+            )
+        else:
+            ds.apply(doc_analyze, ocr=False).pipe_txt_mode(image_writer).dump_md(
+                md_writer,
+                os.path.basename(markdown_path),  # 文件名部分
+                image_dir
+            )
+
+        with open(markdown_path, "r", encoding='utf-8') as f:
+            markdown_content = f.read()
+
+        return markdown_content
+
+pdf_processor = PdfProcessor()
+
+# 使用示例
+if __name__ == "__main__":
+    # pdf_processor = PdfProcessor()
+    print(pdf_processor.process_pdf(
+        "/home/caocaiyu/datamax-service/backend/uploaded_files/fde1daee-e899-4e93-87ff-706234c399c3/20250227132500_5447d25cbf094a3295f9d52d3408a048.pdf"
+    ))

datamax/utils/paddleocr_pdf_operator.py

@@ -8,6 +8,7 @@ from PIL import Image
 from copy import deepcopy
 from datetime import datetime
 
+os.environ['KMP_DUPLICATE_LIB_OK']='True'
 ROOT_DIR: pathlib.Path = pathlib.Path(__file__).parent.parent.parent.resolve()
 sys.path.append(str(ROOT_DIR))
 
@@ -82,7 +83,7 @@ def use_paddleocr(input_files: str, output_files: str, use_gpu: bool = False, gp
         if not os.path.exists(output_files):
             os.makedirs(output_files)
         try:
-            recovery(input_files, output_files, use_gpu, gpu_id)
+            recovery(img_path=input_files, output=output_files, use_gpu=use_gpu, gpu_id=gpu_id)
         except Exception as e:
             raise e
     except Exception as e: