MinerU/magic_pdf/libs/commons.py

import datetime
import json
import os, re, configparser
import subprocess
import time

import boto3
from loguru import logger
from boto3.s3.transfer import TransferConfig
from botocore.config import Config

import fitz # 1.23.9中已经切换到rebase
# import fitz_old as fitz  # 使用1.23.9之前的pymupdf库


def get_delta_time(input_time):
    return round(time.time() - input_time, 2)


def join_path(*args):
    return '/'.join(str(s).rstrip('/') for s in args)


#配置全局的errlog_path，方便demo同步引用
error_log_path = "s3://llm-pdf-text/err_logs/"
# json_dump_path = "s3://pdf_books_temp/json_dump/" # 这条路径仅用于临时本地测试,不能提交到main
json_dump_path = "s3://llm-pdf-text/json_dump/"

# s3_image_save_path = "s3://mllm-raw-media/pdf2md_img/" # 基础库不应该有这些存在的路径，应该在业务代码中定义


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 formatted_time(time_stamp):
    dt_object = datetime.datetime.fromtimestamp(time_stamp)
    output_time = dt_object.strftime("%Y-%m-%d-%H:%M:%S")
    return output_time


def mymax(alist: list):
    if len(alist) == 0:
        return 0  # 空是0， 0*0也是0大小q
    else:
        return max(alist)

def parse_aws_param(profile):
    if isinstance(profile, str):
        # 解析配置文件
        config_file = join_path(os.path.expanduser("~"), ".aws", "config")
        credentials_file = join_path(os.path.expanduser("~"), ".aws", "credentials")
        config = configparser.ConfigParser()
        config.read(credentials_file)
        config.read(config_file)
        # 获取 AWS 账户相关信息
        ak = config.get(profile, "aws_access_key_id")
        sk = config.get(profile, "aws_secret_access_key")
        if profile == "default":
            s3_str = config.get(f"{profile}", "s3")
        else:
            s3_str = config.get(f"profile {profile}", "s3")
        end_match = re.search("endpoint_url[\s]*=[\s]*([^\s\n]+)[\s\n]*$", s3_str, re.MULTILINE)
        if end_match:
            endpoint = end_match.group(1)
        else:
            raise ValueError(f"aws 配置文件中没有找到 endpoint_url")
        style_match = re.search("addressing_style[\s]*=[\s]*([^\s\n]+)[\s\n]*$", s3_str, re.MULTILINE)
        if style_match:
            addressing_style = style_match.group(1)
        else:
            addressing_style = "path"
    elif isinstance(profile, dict):
        ak = profile["ak"]
        sk = profile["sk"]
        endpoint = profile["endpoint"]
        addressing_style = "auto"

    return ak, sk, endpoint, addressing_style


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


def read_file(pdf_path: str, s3_profile):
    if pdf_path.startswith("s3://"):
        ak, sk, end_point, addressing_style = parse_aws_param(s3_profile)
        cli = boto3.client(service_name="s3", aws_access_key_id=ak, aws_secret_access_key=sk, endpoint_url=end_point,
                           config=Config(s3={'addressing_style': addressing_style}, retries={'max_attempts': 10, 'mode': 'standard'}))
        bucket_name, bucket_key = parse_bucket_key(pdf_path)
        res = cli.get_object(Bucket=bucket_name, Key=bucket_key)
        file_content = res["Body"].read()
        return file_content
    else:
        with open(pdf_path, "rb") as f:
            return f.read()


def get_docx_model_output(pdf_model_output, page_id):

    model_output_json = pdf_model_output[page_id]

    return model_output_json


def list_dir(dir_path:str, s3_profile:str):
    """
    列出dir_path下的所有文件
    """
    ret = []
    
    if dir_path.startswith("s3"):
        ak, sk, end_point, addressing_style = parse_aws_param(s3_profile)
        s3info = re.findall(r"s3:\/\/([^\/]+)\/(.*)", dir_path)
        bucket, path = s3info[0][0], s3info[0][1]
        try:
            cli = boto3.client(service_name="s3", aws_access_key_id=ak, aws_secret_access_key=sk, endpoint_url=end_point,
                                            config=Config(s3={'addressing_style': addressing_style}))
            def list_obj_scluster():
                marker = None
                while True:
                    list_kwargs = dict(MaxKeys=1000, Bucket=bucket, Prefix=path)
                    if marker:
                        list_kwargs['Marker'] = marker
                    response = cli.list_objects(**list_kwargs)
                    contents = response.get("Contents", [])
                    yield from contents
                    if not response.get("IsTruncated") or len(contents)==0:
                        break
                    marker = contents[-1]['Key']


            for info in list_obj_scluster():
                file_path = info['Key']
                #size = info['Size']

                if path!="":
                    afile = file_path[len(path):]
                    if afile.endswith(".json"):
                        ret.append(f"s3://{bucket}/{file_path}")
                        
            return ret

        except Exception as e:
            logger.exception(e)
            exit(-1)
    else: #本地的目录，那么扫描本地目录并返会这个目录里的所有jsonl文件
        
        for root, dirs, files in os.walk(dir_path):
            for file in files:
                if file.endswith(".json"):
                    ret.append(join_path(root, file))
        ret.sort()
        return ret

def get_img_s3_client(save_path:str, image_s3_config:str):
    """
    """
    if save_path.startswith("s3://"):  # 放这里是为了最少创建一个s3 client
        ak, sk, end_point, addressing_style = parse_aws_param(image_s3_config)
        img_s3_client = boto3.client(
            service_name="s3",
            aws_access_key_id=ak,
            aws_secret_access_key=sk,
            endpoint_url=end_point,
            config=Config(s3={"addressing_style": addressing_style}, retries={'max_attempts': 5, 'mode': 'standard'}),
        )
    else:
        img_s3_client = None
        
    return img_s3_client

if __name__=="__main__":
    s3_path = "s3://llm-pdf-text/layout_det/scihub/scimag07865000-07865999/10.1007/s10729-011-9175-6.pdf/"
    s3_profile = "langchao"
    ret = list_dir(s3_path, s3_profile)
    print(ret)