Remove unused modules/assets from earlier upload

NaviOCR/data_reader_writer/__init__.py

@@ -1,10 +0,0 @@
-from .base import DataReader, DataWriter
-from .filebase import FileBasedDataReader, FileBasedDataWriter
-from .imagefile import ImageDataWriter
-__all__ = [
-    "DataReader",
-    "DataWriter",
-    "FileBasedDataReader",
-    "FileBasedDataWriter",
-    "ImageDataWriter",
-]

NaviOCR/data_reader_writer/base.py

@@ -1,62 +0,0 @@
-
-from abc import ABC, abstractmethod
-
-
-class DataReader(ABC):
-
-    def read(self, path: str) -> bytes:
-        """Read the file.
-
-        Args:
-            path (str): file path to read
-
-        Returns:
-            bytes: the content of the file
-        """
-        return self.read_at(path)
-
-    @abstractmethod
-    def read_at(self, path: str, offset: int = 0, limit: int = -1) -> bytes:
-        """Read the file at offset and limit.
-
-        Args:
-            path (str): the file path
-            offset (int, optional): the number of bytes skipped. Defaults to 0.
-            limit (int, optional): the length of bytes want to read. Defaults to -1.
-
-        Returns:
-            bytes: the content of the file
-        """
-        pass
-
-
-class DataWriter(ABC):
-    @abstractmethod
-    def write(self, path: str, data: bytes) -> None:
-        """Write the data to the file.
-
-        Args:
-            path (str): the target file where to write
-            data (bytes): the data want to write
-        """
-        pass
-
-    def write_string(self, path: str, data: str) -> None:
-        """Write the data to file, the data will be encoded to bytes.
-
-        Args:
-            path (str): the target file where to write
-            data (str): the data want to write
-        """
-
-        def safe_encode(data: str, method: str):
-            try:
-                bit_data = data.encode(encoding=method, errors='replace')
-                return bit_data, True
-            except:  # noqa
-                return None, False
-        for method in ['utf-8', 'ascii']:
-            bit_data, flag = safe_encode(data, method)
-            if flag:
-                self.write(path, bit_data)
-                break

NaviOCR/data_reader_writer/filebase.py

@@ -1,62 +0,0 @@
-import os
-
-from .base import DataReader, DataWriter
-
-
-class FileBasedDataReader(DataReader):
-    def __init__(self, parent_dir: str = ''):
-        """Initialized with parent_dir.
-
-        Args:
-            parent_dir (str, optional): the parent directory that may be used within methods. Defaults to ''.
-        """
-        self._parent_dir = parent_dir
-
-    def read_at(self, path: str, offset: int = 0, limit: int = -1) -> bytes:
-        """Read at offset and limit.
-
-        Args:
-            path (str): the path of file, if the path is relative path, it will be joined with parent_dir.
-            offset (int, optional): the number of bytes skipped. Defaults to 0.
-            limit (int, optional): the length of bytes want to read. Defaults to -1.
-
-        Returns:
-            bytes: the content of file
-        """
-        fn_path = path
-        if not os.path.isabs(fn_path) and len(self._parent_dir) > 0:
-            fn_path = os.path.join(self._parent_dir, path)
-
-        with open(fn_path, 'rb') as f:
-            f.seek(offset)
-            if limit == -1:
-                return f.read()
-            else:
-                return f.read(limit)
-
-
-class FileBasedDataWriter(DataWriter):
-    def __init__(self, parent_dir: str = '') -> None:
-        """Initialized with parent_dir.
-
-        Args:
-            parent_dir (str, optional): the parent directory that may be used within methods. Defaults to ''.
-        """
-        self._parent_dir = parent_dir
-
-    def write(self, path: str, data: bytes) -> None:
-        """Write file with data.
-
-        Args:
-            path (str): the path of file, if the path is relative path, it will be joined with parent_dir.
-            data (bytes): the data want to write
-        """
-        fn_path = path
-        if not os.path.isabs(fn_path) and len(self._parent_dir) > 0:
-            fn_path = os.path.join(self._parent_dir, path)
-
-        if not os.path.exists(os.path.dirname(fn_path)) and os.path.dirname(fn_path) != "":
-            os.makedirs(os.path.dirname(fn_path), exist_ok=True)
-
-        with open(fn_path, 'wb') as f:
-            f.write(data)

NaviOCR/data_reader_writer/imagefile.py

@@ -1,66 +0,0 @@
-import os
-from io import BytesIO
-from PIL import Image
-from typing import Any, Dict, List
-
-class ImageDataWriter:
-    def __init__(self, parent_dir: str = "", format = "JPEG") -> None:
-
-        self._parent_dir = parent_dir
-        self.images = {} 
-        self.format = format
-
-    def add_image(self, img_bytes, image_name: str) -> None:
-        self.images[image_name] = img_bytes
-
-    def save_image(self, image_name: str) -> None:
-        if image_name not in self.images:
-            raise ValueError(f"Image {image_name} not found in memory.")
-
-        img_bytes = self.images[image_name]
-        file_path = os.path.join(self._parent_dir, image_name)
-        
-        os.makedirs(os.path.dirname(file_path), exist_ok=True)
-
-        with open(file_path, "wb") as f:
-            f.write(img_bytes)
-
-    def save_all_images(self):
-        os.makedirs(self._parent_dir, exist_ok=True)
-        for image_name in self.images:
-            self.save_image(image_name)
-        self.clear_buffer()
-    
-    def clear_buffer(self):
-        self.images.clear()
-
-
-
-def replace_image_path_in_json(
-    contents: Any,
-    mapping: Dict[str, str],
-    sub_type: List[str],
-    key_name: str = "image_path",
-) -> Any:
-    def _replace_value(filename: str) -> str:
-        if filename in mapping:
-            return mapping[filename]
-        raise KeyError(f"Missing mapping for {filename}")
-
-    def _walk(obj: Any) -> Any:
-        if isinstance(obj, dict):
-            return {
-                k: _replace_value(v) if k == key_name and isinstance(v, str) else _walk(v)
-                for k, v in obj.items()
-            }
-
-        if isinstance(obj, list):
-            return [_walk(i) for i in obj]
-
-        return obj
-
-    for i, content in enumerate(contents):
-        if isinstance(content, dict) and content.get("type") in sub_type:
-            contents[i] = _walk(content)
-
-    return contents

NaviOCR/engine.py

@@ -1,161 +0,0 @@
-import json
-import os
-import time
-from loguru import logger
-
-from NaviOCR.tools.pdf_image_tools import convert_pdf_bytes_to_bytes
-from NaviOCR.data_reader_writer import FileBasedDataWriter, ImageDataWriter
-from NaviOCR.tools.draw_bbox import draw_layout_bbox
-from NaviOCR.src.vlm_middle_json_mkcontent import union_make
-from NaviOCR.src.vlm_analyze import doc_analyze
-from NaviOCR.src.vlm_analyze import aio_doc_analyze 
-
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-def prepare_env(output_dir, pdf_file_name):
-    local_md_dir = str(os.path.join(output_dir, pdf_file_name))
-    local_image_dir = os.path.join(str(local_md_dir), "images")
-    os.makedirs(local_image_dir, exist_ok=True)
-    os.makedirs(local_md_dir, exist_ok=True)
-    return local_image_dir, local_md_dir
-
-
-def _prepare_pdf_bytes(pdf_bytes_list, valid_page_ids):
-    result = []
-    for idx, pdf_bytes in enumerate(pdf_bytes_list):
-        valid_single_page_ids = (
-            valid_page_ids[idx]
-            if valid_page_ids and idx < len(valid_page_ids)
-            else None
-        )
-        new_pdf_bytes = convert_pdf_bytes_to_bytes(pdf_bytes, valid_single_page_ids)
-        result.append(new_pdf_bytes)
-    return result
-
-
-def _process_output(
-        pdf_info,
-        pdf_bytes,
-        pdf_file_name,
-        local_md_dir,
-        local_image_dir,
-        md_writer,
-        middle_json,
-        f_draw_layout_bbox=True,
-        f_dump_md=True,
-        f_dump_middle_json=True,
-):
-    if f_draw_layout_bbox:
-        draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir, f"{pdf_file_name}_layout.pdf")
-    image_dir = str(os.path.basename(local_image_dir))
-    if f_dump_md:
-        md_content_str = union_make(pdf_info, image_dir) 
-        md_writer.write_string(
-            f"{pdf_file_name}.md",
-            md_content_str,
-        )
-    if f_dump_middle_json:
-        md_writer.write_string(
-            f"{pdf_file_name}_middle.json",
-            json.dumps(middle_json, ensure_ascii=False, indent=4),
-        )
-    logger.info(f"local output dir is {local_md_dir}")
-
-async def _async_process_vlm(output_dir, pdf_file_names, pdf_bytes_list, **kwargs,):
-    results = []
-    for idx, pdf_bytes in enumerate(pdf_bytes_list):
-        pdf_file_name = pdf_file_names[idx]
-        local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name)
-        image_writer = ImageDataWriter(local_image_dir)
-        vlm_doc_analyze_time = time.time()
-        middle_json = await aio_doc_analyze(
-            pdf_bytes, image_writer=image_writer, **kwargs,
-        )
-        vlm_doc_analyze_time = round(time.time() - vlm_doc_analyze_time, 2)
-        logger.debug(f"doc_analyze cost: {vlm_doc_analyze_time}")
-        
-        pdf_info = middle_json["pdf_info"]
-        md_writer = FileBasedDataWriter(local_md_dir)
-        
-        process_output_time = time.time()
-        _process_output(
-            pdf_info, pdf_bytes, pdf_file_name, local_md_dir, local_image_dir,
-            md_writer, middle_json
-        )
-        image_writer.save_all_images()
-        process_output_time = round(time.time() - process_output_time, 2)
-        logger.debug(f"process_output_time cost: {process_output_time}")
-        results.append(middle_json)
-    return results
-
-
-def _process_vlm(output_dir, pdf_file_names, pdf_bytes_list, **kwargs,):
-    results = []
-    for idx, pdf_bytes in enumerate(pdf_bytes_list):
-        pdf_file_name = pdf_file_names[idx]
-        local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name)
-        image_writer = ImageDataWriter(local_image_dir)
-        
-        vlm_doc_analyze_time = time.time()
-        middle_json = doc_analyze(
-            pdf_bytes, image_writer=image_writer, **kwargs,
-        )
-        vlm_doc_analyze_time = round(time.time() - vlm_doc_analyze_time, 2)
-        logger.debug(f"doc_analyze cost: {vlm_doc_analyze_time}")
-        
-        pdf_info = middle_json["pdf_info"]
-        md_writer = FileBasedDataWriter(local_md_dir)
-        
-        process_output_time =  time.time()
-        _process_output(
-            pdf_info, pdf_bytes, pdf_file_name, local_md_dir, local_image_dir,
-            md_writer, middle_json
-        )
-        image_writer.save_all_images()
-        process_output_time = round(time.time() - process_output_time, 2)
-        logger.debug(f"process_output_time cost: {process_output_time}")
-        
-        results.append(middle_json)
-    return results
-
-
-def inplace_change_page_ids(results, valid_page_ids):
-    for result, valid_page_id in zip(results, valid_page_ids):
-        if valid_page_id is None:
-            continue
-        for page in result.get("pages", []):
-            page_id = page.get("page_id")
-            if (
-                isinstance(page_id, int)
-                and 0 <= page_id < len(valid_page_id)
-            ):
-                page["page_id"] = valid_page_id[page_id]
-    return results
-
-
-async def aio_do_parse( 
-        output_dir,
-        pdf_file_names: list[str],
-        pdf_bytes_list: list[bytes],
-        valid_page_ids: list[list[int] | None],
-        **kwargs,
-):
-    pdf_bytes_list = _prepare_pdf_bytes(pdf_bytes_list, valid_page_ids)
-    Results = await _async_process_vlm(
-        output_dir, pdf_file_names, pdf_bytes_list, **kwargs,
-    )
-    return Results
-
-
-def do_parse(
-        output_dir,
-        pdf_file_names: list[str],
-        pdf_bytes_list: list[bytes],
-        valid_page_ids: list[list[int] | None],
-        **kwargs,
-):
-    pdf_bytes_list = _prepare_pdf_bytes(pdf_bytes_list, valid_page_ids)
-    Results = _process_vlm(
-        output_dir, pdf_file_names, pdf_bytes_list, **kwargs,
-    )
-    return Results

NaviOCR/src/__init__.py

@@ -1 +0,0 @@
-

NaviOCR/src/model_output_to_middle_json.py

@@ -1,37 +0,0 @@
-from tqdm import tqdm
-from NaviOCR.tools.pdf_image_tools import get_page_size
-from NaviOCR.src.vlm_magic_model import MagicModel
-from NaviOCR.tools.cut_image import cut_image_and_table 
-from NaviOCR.tools.enum_class import ContentType
-from NaviOCR.tools.hash_utils import bytes_md5
-from NaviOCR.version import __version__
-
-
-def blocks_to_page_info(page_blocks,image_dict,page,image_writer,page_index) -> dict:
-    scale = image_dict["scale"]
-    page_pil_img = image_dict["img_pil"]
-    page_img_md5 = bytes_md5(page_pil_img.tobytes())
-    width, height = map(int, get_page_size(page))
-    magic_model = MagicModel(page_blocks, width, height)
-    all_spans = magic_model.get_all_spans()
-    for span in all_spans:
-        if span["type"] in [
-            ContentType.IMAGE,
-            ContentType.SEAL,
-            ContentType.CHAR,
-        ]:
-            cut_image_and_table(span,page_pil_img,page_img_md5,page_index,image_writer,scale=scale)
-    page_blocks = magic_model.get_page_blocks()
-    page_info = {"para_blocks": page_blocks,"discarded_blocks": [],"page_size": [width, height],"page_idx": page_index}
-    return page_info
-
-def result_to_middle_json(model_output_blocks_list, images_list, pdf_doc, image_writer):
-    middle_json = {"pdf_info": [],"_backend": "vlm","_version_name": __version__,}
-    for index, page_blocks in enumerate(tqdm(model_output_blocks_list)):
-        page = pdf_doc[index]
-        image_dict = images_list[index]
-        page_info = blocks_to_page_info(page_blocks,image_dict,page,image_writer,index)
-        middle_json['pdf_info'].append(page_info)
-    
-    pdf_doc.close()
-    return middle_json

NaviOCR/src/vlm_analyze.py

@@ -1,63 +0,0 @@
-import time
-from loguru import logger
-from ..data_reader_writer import DataWriter
-from ..tools.enum_class import ImageType
-from .model_output_to_middle_json import result_to_middle_json
-from NaviOCR.vlm_utils.NaviOCR_model import NaviOCRMODEL_SERVICE
-from NaviOCR.tools.pdf_image_tools import load_images_from_pdf
-import NaviOCR.config as CONFIG
-
-def _init_model():
-    backend=CONFIG.BACKEND
-    model_path = CONFIG.model_path
-    predictor = NaviOCRMODEL_SERVICE.get_model(backend, model_path, None)
-    return predictor
-
-async def aio_doc_analyze(
-    pdf_bytes, 
-    predictor = None,
-    image_writer = None,
-):  
-    if predictor is None:
-        predictor = _init_model()
-    load_images_start = time.time()
-    images_list, pdf_doc = load_images_from_pdf(pdf_bytes, image_type=ImageType.PIL, threads=CONFIG.PDF_TOOLS_WORKER_MAX_NUM)
-    images_pil_list = [image_dict["img_pil"] for image_dict in images_list]
-    load_images_time = round(time.time() - load_images_start, 2)
-    logger.debug(f"load images cost: {load_images_time}, speed: {round(load_images_time/len(images_pil_list), 3)} images/s")
-    infer_start = time.time()
-    results = await predictor.aio_batch_two_step_extract(images=images_pil_list)
-    infer_time = round(time.time() - infer_start, 2)
-    logger.debug(f"infer finished, cost: {infer_time}, speed: {round(infer_time / len(results), 3)} page/s")
-    
-    output_start = time.time() 
-    middle_json = result_to_middle_json(results, images_list, pdf_doc, image_writer) 
-    output_time = round(time.time() - output_start, 2)
-    logger.debug(f"output json finished, cost: {output_time}, speed: {round(output_time / len(results), 3)} page/s")
-    return middle_json
-
-
-def doc_analyze(
-    pdf_bytes,
-    image_writer: DataWriter | None,
-    predictor = None,
-):
-    if predictor is None:
-        predictor = _init_model()
-    
-    load_images_start = time.time()
-    images_list, pdf_doc = load_images_from_pdf(pdf_bytes, image_type=ImageType.PIL, threads=CONFIG.PDF_TOOLS_WORKER_MAX_NUM)
-    images_pil_list = [image_dict["img_pil"] for image_dict in images_list]
-    load_images_time = round(time.time() - load_images_start, 2)
-    logger.debug(f"load images cost: {load_images_time}, speed: {round(load_images_time/len(images_pil_list), 3)} images/s")
-
-    infer_start = time.time()
-    results = predictor.batch_two_step_extract(images=images_pil_list)
-    infer_time = round(time.time() - infer_start, 2)
-    logger.debug(f"infer finished, cost: {infer_time}, speed: {round(len(results)/infer_time, 3)} page/s")
-
-    output_start = time.time() 
-    middle_json = result_to_middle_json(results, images_list, pdf_doc, image_writer) 
-    output_time = round(time.time() - output_start, 2)
-    logger.debug(f"output json finished, cost: {output_time}, speed: {round(output_time / len(results), 3)} page/s")
-    return middle_json

NaviOCR/src/vlm_magic_model.py

@@ -1,561 +0,0 @@
-import re
-from typing import Literal
-
-from loguru import logger
-
-from NaviOCR.tools.boxbase import calculate_overlap_area_in_bbox1_area_ratio
-from NaviOCR.tools.enum_class import ContentType, BlockType
-from NaviOCR.tools.guess_suffix_or_lang import guess_language_by_text
-from NaviOCR.tools.magic_model_utils import reduct_overlap, tie_up_category_by_index
-
-
-class MagicModel:
-    def __init__(self, page_blocks: list, width, height):
-        self.page_blocks = page_blocks
-
-        blocks = []
-        self.all_spans = []
-        for index, block_info in enumerate(page_blocks):
-            block_bbox = block_info["bbox"]
-            try:
-                if len(block_bbox) % 2 != 0:
-                    logger.warning(f"Invalid bbox length: {block_bbox}")
-                else:
-                    pts = []
-                    for i in range(0, len(block_bbox), 2):
-                        pts.append(int(block_bbox[i] * width))
-                        pts.append(int(block_bbox[i + 1] * height))
-                    block_bbox = tuple(pts)
-                block_type = block_info["type"]
-                block_content = block_info["content"]
-                block_angle = block_info["angle"]
-            except Exception as e:
-                logger.warning(f"Invalid block format: {block_info}, error: {e}")
-                continue
-            
-            span_type = "unknown"
-            code_block_sub_type = None
-            guess_lang = None
-
-            # 确定类型
-            if block_type in ["list"]:
-                span_type = ContentType.TEXT
-            
-            elif block_type in [
-                "text",
-                "title",
-                "image_caption",
-                "image_footnote",
-                "table_caption",
-                "table_footnote",
-                "code_caption",
-                "ref_text",
-                "phonetic",
-                "header",
-                "footer",
-                "page_number",
-                "aside_text",
-                "page_footnote",
-            ]:
-                span_type = ContentType.TEXT
-            elif block_type in ["image"]:
-                block_type = BlockType.IMAGE_BODY
-                span_type = ContentType.IMAGE
-            elif block_type in ["table"]:
-                block_type = BlockType.TABLE_BODY
-                span_type = ContentType.TABLE
-            elif block_type in ["code", "algorithm"]:
-                guess_lang, block_content = guess_language_by_text(block_content)
-                code_block_sub_type = block_type
-                block_type = BlockType.CODE_BODY
-                span_type = ContentType.TEXT
-            elif block_type in ["equation"]:
-                block_type = BlockType.INTERLINE_EQUATION
-                span_type = ContentType.INTERLINE_EQUATION
-            elif block_type in ['seal']:
-                block_type = BlockType.SEAL
-                span_type = ContentType.SEAL
-            elif block_type in ['char']:
-                block_type = BlockType.CHAR
-                span_type = ContentType.CHAR
-            
-            switch_code_to_algorithm = False
-            
-            if span_type in [ContentType.IMAGE, ContentType.TABLE]:
-                span = {
-                    "bbox": block_bbox,
-                    "type": span_type,
-                }
-                if span_type == ContentType.TABLE:
-                    span["html"] = block_content
-            
-            elif span_type in [ContentType.INTERLINE_EQUATION]:
-                span = {
-                    "bbox": block_bbox,
-                    "type": span_type, 
-                    "content": block_content,
-                }
-            elif span_type in [ContentType.SEAL]:
-                span = {
-                    "bbox": block_bbox,
-                    "type": span_type, 
-                    "content": block_content
-                }
-            elif span_type in [ContentType.CHAR]:
-                span = {
-                    "bbox": block_bbox,
-                    "type": span_type, 
-                    "content": block_content
-                }
-            else:
-                if block_type == "title" and block_content:
-                    block_content = re.sub(r"\n\s*", " ", block_content).strip()
-                span = {
-                    "bbox": block_bbox,
-                    "type": span_type,
-                    "content": block_content,
-                }
-
-            if isinstance(span, dict) and "bbox" in span:
-                self.all_spans.append(span)
-                spans = [span]
-            elif isinstance(span, list):
-                self.all_spans.extend(span)
-                spans = span
-            else:
-                raise ValueError(f"Invalid span type: {span_type}, expected dict or list, got {type(span)}")
-
-            if block_type in [BlockType.CODE_BODY]:
-                if switch_code_to_algorithm and code_block_sub_type == "code":
-                    code_block_sub_type = "algorithm"
-                line = {"bbox": block_bbox, "spans": spans, "extra": {"type": code_block_sub_type, "guess_lang": guess_lang}}
-            else:
-                line = {"bbox": block_bbox, "spans": spans}
-
-            blocks.append(
-                {
-                    "bbox": block_bbox,
-                    "type": block_type,
-                    "angle": block_angle,
-                    "lines": [line],
-                    "index": index,
-                }
-            ) 
-
-        # 收集各式的blocks
-        self.image_blocks = []
-        self.table_blocks = []
-        self.interline_equation_blocks = []
-        self.text_blocks = []
-        self.title_blocks = []
-        self.code_blocks = []
-        self.discarded_blocks = []
-        self.ref_text_blocks = []
-        self.phonetic_blocks = []
-        self.list_blocks = []
-        self.seal_blocks = []
-        self.char_blocks = []
-        for block in blocks:
-            if block["type"] in [BlockType.IMAGE_BODY, BlockType.IMAGE_CAPTION, BlockType.IMAGE_FOOTNOTE]:
-                self.image_blocks.append(block)
-            elif block["type"] in [BlockType.TABLE_BODY, BlockType.TABLE_CAPTION, BlockType.TABLE_FOOTNOTE]:
-                self.table_blocks.append(block)
-            elif block["type"] in [BlockType.CODE_BODY, BlockType.CODE_CAPTION]:
-                self.code_blocks.append(block)
-            elif block["type"] == BlockType.INTERLINE_EQUATION:
-                self.interline_equation_blocks.append(block)
-            elif block["type"] == BlockType.TEXT:
-                self.text_blocks.append(block) 
-            elif block["type"] == BlockType.TITLE:
-                self.title_blocks.append(block)
-            elif block["type"] in [BlockType.REF_TEXT]:
-                self.ref_text_blocks.append(block)
-            elif block["type"] in [BlockType.PHONETIC]:
-                self.phonetic_blocks.append(block)
-            elif block["type"] in [BlockType.HEADER, BlockType.FOOTER, BlockType.PAGE_NUMBER, BlockType.ASIDE_TEXT, BlockType.PAGE_FOOTNOTE]:
-                self.discarded_blocks.append(block)
-            elif block["type"] == BlockType.LIST:
-                self.list_blocks.append(block)
-            elif block["type"] == BlockType.SEAL:
-                self.seal_blocks.append(block)
-            elif block["type"] == BlockType.CHAR:
-                self.char_blocks.append(block)
-            else:
-                continue
-        
-
-        self.list_blocks, self.text_blocks, self.ref_text_blocks = fix_list_blocks(self.list_blocks, self.text_blocks, self.ref_text_blocks)
-        self.image_blocks, not_include_image_blocks = fix_two_layer_blocks(self.image_blocks, BlockType.IMAGE)
-        self.table_blocks, not_include_table_blocks = fix_two_layer_blocks(self.table_blocks, BlockType.TABLE)
-        self.code_blocks, not_include_code_blocks = fix_two_layer_blocks(self.code_blocks, BlockType.CODE)
-        for code_block in self.code_blocks:
-            for block in code_block['blocks']:
-                if block['type'] == BlockType.CODE_BODY:
-                    if len(block["lines"]) > 0:
-                        line = block["lines"][0]
-                        code_block["sub_type"] = line["extra"]["type"]
-                        if code_block["sub_type"] in ["code"]:
-                            code_block["guess_lang"] = line["extra"]["guess_lang"]
-                        del line["extra"]
-                    else:
-                        code_block["sub_type"] = "code"
-                        code_block["guess_lang"] = "txt"
-
-
-        for block in not_include_image_blocks + not_include_table_blocks + not_include_code_blocks:
-            block["type"] = BlockType.TEXT
-            self.text_blocks.append(block)
-
-
-    def get_list_blocks(self):
-
-        return self.list_blocks
-
-    def get_image_blocks(self):
-
-        return self.image_blocks
-
-    def get_table_blocks(self):
-
-        return self.table_blocks
-
-    def get_code_blocks(self):
-
-        return self.code_blocks
-
-    def get_ref_text_blocks(self):
-
-        return self.ref_text_blocks
-
-    def get_phonetic_blocks(self):
-
-        return self.phonetic_blocks
-
-    def get_title_blocks(self):
-
-        return self.title_blocks
-
-    def get_text_blocks(self):
-
-        return self.text_blocks
-
-    def get_interline_equation_blocks(self):
-
-        return self.interline_equation_blocks
-
-    def get_discarded_blocks(self):
-
-        return self.discarded_blocks
-
-    def get_all_spans(self):
-        return self.all_spans
-
-    def get_content_blocks(self, content_blocks):
-        new_content_blocks = []
-        for inter_block in  content_blocks:
-            contents = extract_contents(inter_block.get('lines', []), 'content')
-            if len(contents) > 0:
-                new_content_blocks.append({
-                    "type":inter_block.get("type", "text"),
-                    "angle":inter_block.get("angle", 0),
-                    "index":inter_block.get('index', -1),
-                    "pos":inter_block.get('bbox', []),
-                    "text":contents[0],
-                })
-        return new_content_blocks
-
-    def get_seal_blocks(self):
-
-        return self.seal_blocks
-
-
-    def get_char_blocks(self):
-
-        return self.char_blocks
-
-    def get_page_blocks(self, pass_blocks=False):
-        image_blocks = self.get_image_blocks()
-        table_blocks = self.get_table_blocks()
-        title_blocks = self.get_title_blocks()
-        discarded_blocks = self.get_discarded_blocks()
-        code_blocks = self.get_code_blocks()
-        ref_text_blocks = self.get_ref_text_blocks()
-        phonetic_blocks = self.get_phonetic_blocks()
-        list_blocks = self.get_list_blocks()
-        text_blocks = self.get_text_blocks()
-        interline_equation_blocks = self.get_interline_equation_blocks()
-        seal_blocks = self.get_seal_blocks()
-        char_blocks = self.get_char_blocks()
-
-        if pass_blocks:
-            page_blocks = [
-                *table_blocks,
-                *code_blocks,
-                *ref_text_blocks,
-                *phonetic_blocks,
-                *title_blocks,
-                *text_blocks,
-                *interline_equation_blocks,
-                *discarded_blocks,
-            ]
-        else:
-            page_blocks = [
-                *image_blocks,
-                *table_blocks,
-                *code_blocks,
-                *ref_text_blocks,
-                *phonetic_blocks,
-                *title_blocks,
-                *text_blocks,
-                *interline_equation_blocks,
-                *list_blocks,
-                *discarded_blocks,
-                *seal_blocks,
-                *char_blocks,
-            ]
-        page_blocks.sort(key=lambda x: x["index"])
-        for i, block in enumerate(page_blocks):
-            block["index"] = i
-        return page_blocks
-        
-                
-
-
-def extract_contents(data, key):
-    image_paths = []
-    def _walk(obj):
-        if isinstance(obj, dict):
-            if key in obj:
-                image_paths.append(obj[key])
-            for v in obj.values():
-                _walk(v)
-        elif isinstance(obj, list):
-            for item in obj:
-                _walk(item)
-    _walk(data)
-    return image_paths
-
-
-def __tie_up_category_by_index(blocks, subject_block_type, object_block_type):
-    """基于index的主客体关联包装函数"""
-    # 定义获取主体和客体对象的函数
-    def get_subjects():
-        return reduct_overlap(
-            list(
-                map(
-                    lambda x: {"bbox": x["bbox"], "lines": x["lines"], "index": x["index"], "angle": x["angle"]},
-                    filter(
-                        lambda x: x["type"] == subject_block_type,
-                        blocks,
-                    ),
-                )
-            )
-        )
-
-    def get_objects():
-        return reduct_overlap(
-            list(
-                map(
-                    lambda x: {"bbox": x["bbox"], "lines": x["lines"], "index": x["index"], "angle": x["angle"]},
-                    filter(
-                        lambda x: x["type"] == object_block_type,
-                        blocks,
-                    ),
-                )
-            )
-        )
-
-    # 调用通用方法
-    return tie_up_category_by_index(
-        get_subjects,
-        get_objects
-    )
-
-
-def get_type_blocks(blocks, block_type: Literal["image", "table", "code"]):
-    with_captions = __tie_up_category_by_index(blocks, f"{block_type}_body", f"{block_type}_caption")
-    with_footnotes = __tie_up_category_by_index(blocks, f"{block_type}_body", f"{block_type}_footnote")
-    ret = []
-    for v in with_captions:
-        record = {
-            f"{block_type}_body": v["sub_bbox"],
-            f"{block_type}_caption_list": v["obj_bboxes"],
-        }
-        filter_idx = v["sub_idx"]
-        d = next(filter(lambda x: x["sub_idx"] == filter_idx, with_footnotes))
-        record[f"{block_type}_footnote_list"] = d["obj_bboxes"]
-        ret.append(record)
-    return ret
-
-
-def fix_two_layer_blocks(blocks, fix_type: Literal["image", "table", "code"]):
-    need_fix_blocks = get_type_blocks(blocks, fix_type)
-    fixed_blocks = []
-    not_include_blocks = []
-    processed_indices = set()
-
-    # 特殊处理表格类型，确保标题在表格前，注脚在表格后
-    if fix_type in ["table", "image"]:
-        # 收集所有不合适的caption和footnote
-        misplaced_captions = []  # 存储(caption, 原始block索引)
-        misplaced_footnotes = []  # 存储(footnote, 原始block索引)
-
-        # 第一步：移除不符合位置要求的footnote
-        for block_idx, block in enumerate(need_fix_blocks):
-            body = block[f"{fix_type}_body"]
-            body_index = body["index"]
-
-            # 检查footnote应在body后或同位置
-            valid_footnotes = []
-            for footnote in block[f"{fix_type}_footnote_list"]:
-                if footnote["index"] >= body_index:
-                    valid_footnotes.append(footnote)
-                else:
-                    misplaced_footnotes.append((footnote, block_idx))
-            block[f"{fix_type}_footnote_list"] = valid_footnotes
-
-        # 第三步：重新分配不合规的footnote到合适的body
-        for footnote, original_block_idx in misplaced_footnotes:
-            footnote_index = footnote["index"]
-            best_block_idx = None
-            min_distance = float('inf')
-
-            # 寻找索引小于等于footnote_index的最近body
-            for idx, block in enumerate(need_fix_blocks):
-                body_index = block[f"{fix_type}_body"]["index"]
-                if body_index <= footnote_index and idx != original_block_idx:
-                    distance = footnote_index - body_index
-                    if distance < min_distance:
-                        min_distance = distance
-                        best_block_idx = idx
-
-            if best_block_idx is not None:
-                # 找到合适的body，添加到对应block的footnote_list
-                need_fix_blocks[best_block_idx][f"{fix_type}_footnote_list"].append(footnote)
-            else:
-                # 没找到合适的body，作为普通block处理
-                not_include_blocks.append(footnote)
-
-        # 第四步:将每个block的caption_list和footnote_list中不连续index的元素提出来作为普通block处理
-        for block in need_fix_blocks:
-            caption_list = block[f"{fix_type}_caption_list"]
-            footnote_list = block[f"{fix_type}_footnote_list"]
-            body_index = block[f"{fix_type}_body"]["index"]
-
-            # 处理caption_list (从body往前看,caption在body之前)
-            if caption_list:
-                # 按index降序排列,从最接近body的开始检查
-                caption_list.sort(key=lambda x: x["index"], reverse=True)
-                filtered_captions = [caption_list[0]]
-                for i in range(1, len(caption_list)):
-                    prev_index = caption_list[i - 1]["index"]
-                    curr_index = caption_list[i]["index"]
-
-                    # 检查是否连续
-                    if curr_index == prev_index - 1:
-                        filtered_captions.append(caption_list[i])
-                    else:
-                        # 检查gap中是否只有body_index
-                        gap_indices = set(range(curr_index + 1, prev_index))
-                        if gap_indices == {body_index}:
-                            # gap中只有body_index,不算真正的gap
-                            filtered_captions.append(caption_list[i])
-                        else:
-                            # 出现真正的gap,后续所有caption都作为普通block
-                            not_include_blocks.extend(caption_list[i:])
-                            break
-                # 恢复升序
-                filtered_captions.reverse()
-                block[f"{fix_type}_caption_list"] = filtered_captions
-
-            # 处理footnote_list (从body往后看,footnote在body之后)
-            if footnote_list:
-                # 按index升序排列,从最接近body的开始检查
-                footnote_list.sort(key=lambda x: x["index"])
-                filtered_footnotes = [footnote_list[0]]
-                for i in range(1, len(footnote_list)):
-                    # 检查是否与前一个footnote连续
-                    if footnote_list[i]["index"] == footnote_list[i - 1]["index"] + 1:
-                        filtered_footnotes.append(footnote_list[i])
-                    else:
-                        # 出现gap,后续所有footnote都作为普通block
-                        not_include_blocks.extend(footnote_list[i:])
-                        break
-                block[f"{fix_type}_footnote_list"] = filtered_footnotes
-
-    # 构建两层结构blocks
-    for block in need_fix_blocks:
-        body = block[f"{fix_type}_body"]
-        caption_list = block[f"{fix_type}_caption_list"]
-        footnote_list = block[f"{fix_type}_footnote_list"]
-
-        body["type"] = f"{fix_type}_body"
-        for caption in caption_list:
-            caption["type"] = f"{fix_type}_caption"
-            processed_indices.add(caption["index"])
-        for footnote in footnote_list:
-            footnote["type"] = f"{fix_type}_footnote"
-            processed_indices.add(footnote["index"])
-
-        processed_indices.add(body["index"])
-
-        two_layer_block = {
-            "type": fix_type,
-            "bbox": body["bbox"],
-            "blocks": [body],
-            "index": body["index"],
-        }
-        two_layer_block["blocks"].extend([*caption_list, *footnote_list])
-        # 对blocks按index排序
-        two_layer_block["blocks"].sort(key=lambda x: x["index"])
-
-        fixed_blocks.append(two_layer_block)
-
-    # 添加未处理的blocks
-    for block in blocks:
-        block.pop("type", None)
-        if block["index"] not in processed_indices and block not in not_include_blocks:
-            not_include_blocks.append(block)
-
-    return fixed_blocks, not_include_blocks
-
-
-def fix_list_blocks(list_blocks, text_blocks, ref_text_blocks):
-    for list_block in list_blocks:
-        list_block["blocks"] = []
-        if "lines" in list_block:
-            del list_block["lines"]
-
-    temp_text_blocks = text_blocks + ref_text_blocks
-    need_remove_blocks = []
-    for block in temp_text_blocks:
-        for list_block in list_blocks:
-            if calculate_overlap_area_in_bbox1_area_ratio(block["bbox"], list_block["bbox"]) >= 0.8:
-                list_block["blocks"].append(block)
-                need_remove_blocks.append(block)
-                break
-    
-
-    for block in need_remove_blocks:
-        if block in text_blocks:
-            text_blocks.remove(block)
-        elif block in ref_text_blocks:
-            ref_text_blocks.remove(block)
-
-    # 移除blocks为空的list_block
-    list_blocks = [lb for lb in list_blocks if lb["blocks"]]
-
-    for list_block in list_blocks:
-        # 统计list_block["blocks"]中所有block的type，用众数作为list_block的sub_type
-        type_count = {}
-        for sub_block in list_block["blocks"]:
-            sub_block_type = sub_block["type"]
-            if sub_block_type not in type_count:
-                type_count[sub_block_type] = 0
-            type_count[sub_block_type] += 1
-
-        if type_count:
-            list_block["sub_type"] = max(type_count, key=type_count.get)
-        else:
-            list_block["sub_type"] = "unknown"
-    return list_blocks, text_blocks, ref_text_blocks

NaviOCR/src/vlm_middle_json_mkcontent.py

@@ -1,176 +0,0 @@
-import os
-
-from loguru import logger
-
-from NaviOCR.tools.char_utils import full_to_half_exclude_marks, is_hyphen_at_line_end
-from NaviOCR.tools.config_reader import get_latex_delimiter_config
-from NaviOCR.tools.enum_class import BlockType, ContentType
-from NaviOCR.tools.language import detect_lang
-
-
-def merge_para_with_text(para_block):
-    block_text = ''
-    for line in para_block['lines']:
-        for span in line['spans']:
-            if span['type'] in [ContentType.TEXT]:
-                try:
-                    span['content'] = full_to_half_exclude_marks(span['content'])
-                except:
-                    continue
-                block_text += span['content']
-    block_lang = detect_lang(block_text)
-
-    para_text = ''
-    for i, line in enumerate(para_block['lines']):
-        for j, span in enumerate(line['spans']):
-            span_type = span['type']
-            content = ''
-            if span_type == ContentType.TEXT:
-                content = span['content']
-            elif span_type == ContentType.INLINE_EQUATION:
-                content = span['content']
-            elif span_type == ContentType.INTERLINE_EQUATION:
-                content = span['content']
-            try:
-                content = content.strip()
-            except:
-                continue
-            if content:
-                if span_type == ContentType.INTERLINE_EQUATION:
-                    para_text += content
-                    continue
-                cjk_langs = {'zh', 'ja', 'ko'}
-                is_last_span = j == len(line['spans']) - 1
-
-                if block_lang in cjk_langs:  
-                    if is_last_span and span_type != ContentType.INLINE_EQUATION:
-                        para_text += content
-                    else:
-                        para_text += f'{content} '
-                else:
-                    if span_type in [ContentType.TEXT, ContentType.INLINE_EQUATION]:
-                        if (
-                                is_last_span
-                                and span_type == ContentType.TEXT
-                                and is_hyphen_at_line_end(content)
-                        ):
-                            if (
-                                    i+1 < len(para_block['lines'])
-                                    and para_block['lines'][i + 1].get('spans')
-                                    and para_block['lines'][i + 1]['spans'][0].get('type') == ContentType.TEXT
-                                    and para_block['lines'][i + 1]['spans'][0].get('content', '')
-                                    and para_block['lines'][i + 1]['spans'][0]['content'][0].islower()
-                            ):
-                                para_text += content[:-1]
-                            else:
-                                para_text += content
-                        else:
-                            para_text += f'{content} '
-    return para_text
-
-
-def mk_blocks_to_markdown(para_blocks, img_buket_path=''):
-    page_markdown = []
-    for para_block in para_blocks:
-        para_text = ''
-        para_type = para_block['type']
-        if para_type in [BlockType.TEXT, BlockType.INTERLINE_EQUATION, BlockType.PHONETIC, BlockType.REF_TEXT, BlockType.FOOTER, BlockType.HEADER]:
-            para_text = merge_para_with_text(para_block)
-        elif para_type == BlockType.LIST:
-            for block in para_block['blocks']:
-                item_text = merge_para_with_text(block)
-                para_text += f"{item_text}  \n"
-        elif para_type == BlockType.TITLE:
-            title_level = get_title_level(para_block)
-            para_text = f'{"#" * title_level} {merge_para_with_text(para_block)}'
-        elif para_type == BlockType.IMAGE:
-            has_image_footnote = any(block['type'] == BlockType.IMAGE_FOOTNOTE for block in para_block['blocks'])
-            if has_image_footnote:
-                for block in para_block['blocks']:  # 1st.拼image_caption
-                    if block['type'] == BlockType.IMAGE_CAPTION:
-                        para_text += merge_para_with_text(block) + '  \n'
-                for block in para_block['blocks']:  # 2nd.拼image_body
-                    if block['type'] == BlockType.IMAGE_BODY:
-                        for line in block['lines']:
-                            for span in line['spans']:
-                                if span['type'] == ContentType.IMAGE:
-                                    if span.get('image_path', ''):
-                                        para_text += f"![]({img_buket_path}/{span['image_path']})"
-                for block in para_block['blocks']:  # 3rd.拼image_footnote
-                    if block['type'] == BlockType.IMAGE_FOOTNOTE:
-                        para_text += '  \n' + merge_para_with_text(block)
-            else:
-                for block in para_block['blocks']: 
-                    if block['type'] == BlockType.IMAGE_BODY:
-                        for line in block['lines']:
-                            for span in line['spans']:
-                                if span['type'] == ContentType.IMAGE:
-                                    if span.get('image_path', ''):
-                                        para_text += f"![]({img_buket_path}/{span['image_path']})"
-                for block in para_block['blocks']:
-                    if block['type'] == BlockType.IMAGE_CAPTION:
-                        para_text += '  \n' + merge_para_with_text(block)
-        elif para_type == BlockType.TABLE:
-            for block in para_block['blocks']:
-                if block['type'] == BlockType.TABLE_CAPTION:
-                    para_text += merge_para_with_text(block) + '  \n'
-            for block in para_block['blocks']:
-                if block['type'] == BlockType.TABLE_BODY:
-                    for line in block['lines']:
-                        for span in line['spans']:
-                            if span['type'] == ContentType.TABLE:
-                                if span.get('html', ''):
-                                    para_text += f"\n{span['html']}\n"
-                                elif span.get('image_path', ''):
-                                    para_text += f"![]({img_buket_path}/{span['image_path']})"
-            for block in para_block['blocks']:
-                if block['type'] == BlockType.TABLE_FOOTNOTE:
-                    para_text += '\n' + merge_para_with_text(block) + '  '
-        elif para_type == BlockType.CODE:
-            sub_type = para_block["sub_type"]
-            for block in para_block['blocks']: 
-                if block['type'] == BlockType.CODE_CAPTION:
-                    para_text += merge_para_with_text(block) + '  \n'
-            for block in para_block['blocks']: 
-                if block['type'] == BlockType.CODE_BODY:
-                    if sub_type == BlockType.CODE:
-                        guess_lang = para_block["guess_lang"]
-                        para_text += f"```{guess_lang}\n{merge_para_with_text(block)}\n```"
-                    elif sub_type == BlockType.ALGORITHM:
-                        para_text += merge_para_with_text(block)
-        elif para_type == BlockType.CHAR: 
-            for line in para_block['lines']:
-                spans = line['spans']
-                for span in spans:
-                    content = span['content']
-                    if len(content) <= 5:
-                        continue
-                    else:
-                        para_text += f"\n{content}\n"
-        if para_text.strip() == '':
-            continue
-        else:
-            page_markdown.append(para_text.strip())
-
-    return page_markdown
-
-
-def union_make(pdf_info_dict: list,img_buket_path: str = '', ): 
-    output_content = []
-    for page_info in pdf_info_dict:
-        paras_of_layout = page_info.get('para_blocks')
-
-        if not paras_of_layout:
-            continue
-        page_markdown = mk_blocks_to_markdown(paras_of_layout, img_buket_path)
-        output_content.extend(page_markdown)
-
-    return '\n\n'.join(output_content)
-
-def get_title_level(block):
-    title_level = block.get('level', 1)
-    if title_level > 4:
-        title_level = 4
-    elif title_level < 1:
-        title_level = 0
-    return title_level

NaviOCR/tools/boxbase.py

@@ -1,224 +0,0 @@
-import math
-
-def normalize_bbox(bbox):
-    if isinstance(bbox[0], (int, float)) and len(bbox) == 4:
-        x0, y0, x1, y1 = bbox
-        return x0, y0, x1, y1
-    if isinstance(bbox[0], (int, float)) and len(bbox) % 2 == 0:
-        xs = bbox[0::2]
-        ys = bbox[1::2]
-        return min(xs), min(ys), max(xs), max(ys)
-    xs = [p[0] for p in bbox]
-    ys = [p[1] for p in bbox]
-    return min(xs), min(ys), max(xs), max(ys)
-
-def is_in(box1, box2) -> bool:
-    x0_1, y0_1, x1_1, y1_1 = normalize_bbox(box1)
-    x0_2, y0_2, x1_2, y1_2 = normalize_bbox(box2)
-    return (
-        x0_1 >= x0_2
-        and y0_1 >= y0_2
-        and x1_1 <= x1_2
-        and y1_1 <= y1_2
-    )
-
-def bbox_relative_pos(bbox1, bbox2):
-    """判断两个矩形框的相对位置关系.
-
-    Args:
-        bbox1: 一个四元组，表示第一个矩形框的左上角和右下角的坐标，格式为(x1, y1, x1b, y1b)
-        bbox2: 一个四元组，表示第二个矩形框的左上角和右下角的坐标，格式为(x2, y2, x2b, y2b)
-
-    Returns:
-        一个四元组，表示矩形框1相对于矩形框2的位置关系，格式为(left, right, bottom, top)
-        其中，left表示矩形框1是否在矩形框2的左侧，right表示矩形框1是否在矩形框2的右侧，
-        bottom表示矩形框1是否在矩形框2的下方，top表示矩形框1是否在矩形框2的上方
-    """
-    x1, y1, x1b, y1b = bbox1
-    x2, y2, x2b, y2b = bbox2
-
-    left = x2b < x1
-    right = x1b < x2
-    bottom = y2b < y1
-    top = y1b < y2
-    return left, right, bottom, top
-
-
-def bbox_distance(bbox1, bbox2):
-    """计算两个矩形框的距离。
-    Args:
-        bbox1 (tuple): 第一个矩形框的坐标，格式为 (x1, y1, x2, y2)，其中 (x1, y1) 为左上角坐标，(x2, y2) 为右下角坐标。
-        bbox2 (tuple): 第二个矩形框的坐标，格式为 (x1, y1, x2, y2)，其中 (x1, y1) 为左上角坐标，(x2, y2) 为右下角坐标。
-    Returns:
-        float: 矩形框之间的距离。
-    """
-
-    def dist(point1, point2):
-        return math.sqrt((point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2)
-
-    x1, y1, x1b, y1b = bbox1
-    x2, y2, x2b, y2b = bbox2
-
-    left, right, bottom, top = bbox_relative_pos(bbox1, bbox2)
-
-    if top and left:
-        return dist((x1, y1b), (x2b, y2))
-    elif left and bottom:
-        return dist((x1, y1), (x2b, y2b))
-    elif bottom and right:
-        return dist((x1b, y1), (x2, y2b))
-    elif right and top:
-        return dist((x1b, y1b), (x2, y2))
-    elif left:
-        return x1 - x2b
-    elif right:
-        return x2 - x1b
-    elif bottom:
-        return y1 - y2b
-    elif top:
-        return y2 - y1b
-    return 0.0
-
-
-def _bbox_center(bbox):
-    if len(bbox) < 4 or len(bbox) % 2 != 0:
-        raise ValueError("bbox must contain an even number of coordinates (>=4)")
-    xs = bbox[0::2]
-    ys = bbox[1::2]
-    n = len(xs)
-    return sum(xs) / n, sum(ys) / n
-
-
-def bbox_center_distance(bbox1, bbox2):
-    c1x, c1y = _bbox_center(bbox1)
-    c2x, c2y = _bbox_center(bbox2)
-    return math.sqrt((c1x - c2x) ** 2 + (c1y - c2y) ** 2)
-
-
-def get_minbox_if_overlap_by_ratio(bbox1, bbox2, ratio):
-    """通过calculate_overlap_area_2_minbox_area_ratio计算两个bbox重叠的面积占最小面积的box的比例
-    如果比例大于ratio，则返回小的那个bbox, 否则返回None."""
-    x1_min, y1_min, x1_max, y1_max = bbox1
-    x2_min, y2_min, x2_max, y2_max = bbox2
-    area1 = (x1_max - x1_min) * (y1_max - y1_min)
-    area2 = (x2_max - x2_min) * (y2_max - y2_min)
-    overlap_ratio = calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2)
-    if overlap_ratio > ratio:
-        if area1 <= area2:
-            return bbox1
-        else:
-            return bbox2
-    else:
-        return None
-
-
-def calculate_overlap_area_2_minbox_area_ratio(bbox1, bbox2):
-    """计算box1和box2的重叠面积占最小面积的box的比例."""
-    # Determine the coordinates of the intersection rectangle
-    x_left = max(bbox1[0], bbox2[0])
-    y_top = max(bbox1[1], bbox2[1])
-    x_right = min(bbox1[2], bbox2[2])
-    y_bottom = min(bbox1[3], bbox2[3])
-
-    if x_right < x_left or y_bottom < y_top:
-        return 0.0
-
-    # The area of overlap area
-    intersection_area = (x_right - x_left) * (y_bottom - y_top)
-    min_box_area = min([(bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1]),
-                        (bbox2[3] - bbox2[1]) * (bbox2[2] - bbox2[0])])
-    if min_box_area == 0:
-        return 0
-    else:
-        return intersection_area / min_box_area
-
-
-def calculate_iou(bbox1, bbox2):
-    """计算两个边界框的交并比(IOU)。
-
-    Args:
-        bbox1 (list[float]): 第一个边界框的坐标，格式为 [x1, y1, x2, y2]，其中 (x1, y1) 为左上角坐标，(x2, y2) 为右下角坐标。
-        bbox2 (list[float]): 第二个边界框的坐标，格式与 `bbox1` 相同。
-
-    Returns:
-        float: 两个边界框的交并比(IOU)，取值范围为 [0, 1]。
-    """
-    # Determine the coordinates of the intersection rectangle
-    x_left = max(bbox1[0], bbox2[0])
-    y_top = max(bbox1[1], bbox2[1])
-    x_right = min(bbox1[2], bbox2[2])
-    y_bottom = min(bbox1[3], bbox2[3])
-
-    if x_right < x_left or y_bottom < y_top:
-        return 0.0
-
-    # The area of overlap area
-    intersection_area = (x_right - x_left) * (y_bottom - y_top)
-
-    # The area of both rectangles
-    bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
-    bbox2_area = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
-
-    if any([bbox1_area == 0, bbox2_area == 0]):
-        return 0
-
-    # Compute the intersection over union by taking the intersection area
-    # and dividing it by the sum of both areas minus the intersection area
-    iou = intersection_area / float(bbox1_area + bbox2_area - intersection_area)
-
-    return iou
-
-
-def calculate_overlap_area_in_bbox1_area_ratio(bbox1, bbox2):
-    """计算box1和box2的重叠面积占bbox1的比例."""
-    # Determine the coordinates of the intersection rectangle
-    x_left = max(bbox1[0], bbox2[0])
-    y_top = max(bbox1[1], bbox2[1])
-    x_right = min(bbox1[2], bbox2[2])
-    y_bottom = min(bbox1[3], bbox2[3])
-
-    if x_right < x_left or y_bottom < y_top:
-        return 0.0
-
-    # The area of overlap area
-    intersection_area = (x_right - x_left) * (y_bottom - y_top)
-    bbox1_area = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
-    if bbox1_area == 0:
-        return 0
-    else:
-        return intersection_area / bbox1_area
-
-
-def calculate_vertical_projection_overlap_ratio(block1, block2):
-    """
-    Calculate the proportion of the x-axis covered by the vertical projection of two blocks.
-
-    Args:
-        block1 (tuple): Coordinates of the first block (x0, y0, x1, y1).
-        block2 (tuple): Coordinates of the second block (x0, y0, x1, y1).
-
-    Returns:
-        float: The proportion of the x-axis covered by the vertical projection of the two blocks.
-    """
-    x0_1, _, x1_1, _ = block1
-    x0_2, _, x1_2, _ = block2
-
-    # Calculate the intersection of the x-coordinates
-    x_left = max(x0_1, x0_2)
-    x_right = min(x1_1, x1_2)
-
-    if x_right < x_left:
-        return 0.0
-
-    # Length of the intersection
-    intersection_length = x_right - x_left
-
-    # Length of the x-axis projection of the first block
-    block1_length = x1_1 - x0_1
-
-    if block1_length == 0:
-        return 0.0
-
-    # Proportion of the x-axis covered by the intersection
-    # logger.info(f"intersection_length: {intersection_length}, block1_length: {block1_length}")
-    return intersection_length / block1_length

NaviOCR/tools/char_utils.py

@@ -1,55 +0,0 @@
-#  Copyright (c) Opendatalab. All rights reserved.
-import re
-
-
-def is_hyphen_at_line_end(line):
-    """Check if a line ends with one or more letters followed by a hyphen.
-
-    Args:
-    line (str): The line of text to check.
-
-    Returns:
-    bool: True if the line ends with one or more letters followed by a hyphen, False otherwise.
-    """
-    # Use regex to check if the line ends with one or more letters followed by a hyphen
-    return bool(re.search(r'[A-Za-z]+-\s*$', line))
-
-
-def full_to_half_exclude_marks(text: str) -> str:
-    """Convert full-width characters to half-width characters using code point manipulation.
-
-    Args:
-        text: String containing full-width characters
-
-    Returns:
-        String with full-width characters converted to half-width
-    """
-    result = []
-    for char in text:
-        code = ord(char)
-        # Full-width letters and numbers (FF21-FF3A for A-Z, FF41-FF5A for a-z, FF10-FF19 for 0-9)
-        if (0xFF21 <= code <= 0xFF3A) or (0xFF41 <= code <= 0xFF5A) or (0xFF10 <= code <= 0xFF19):
-            result.append(chr(code - 0xFEE0))  # Shift to ASCII range
-        else:
-            result.append(char)
-    return ''.join(result)
-
-
-def full_to_half(text: str) -> str:
-    """Convert full-width characters to half-width characters using code point manipulation.
-
-    Args:
-        text: String containing full-width characters
-
-    Returns:
-        String with full-width characters converted to half-width
-    """
-    result = []
-    for char in text:
-        code = ord(char)
-        # Full-width letters, numbers and punctuation (FF01-FF5E)
-        if 0xFF01 <= code <= 0xFF5E:
-            result.append(chr(code - 0xFEE0))  # Shift to ASCII range
-        else:
-            result.append(char)
-    return ''.join(result)

NaviOCR/tools/check_sys_env.py

@@ -1,38 +0,0 @@
-
-import platform
-
-from packaging import version
-
-
-def is_windows_environment() -> bool:
-    return platform.system() == "Windows"
-
-
-# Detect if the current environment is a Mac computer
-def is_mac_environment() -> bool:
-    return platform.system() == "Darwin"
-
-
-def is_linux_environment() -> bool:
-    return platform.system() == "Linux"
-
-
-# Detect if CPU is Apple Silicon architecture
-def is_apple_silicon_cpu() -> bool:
-    return platform.machine() in ["arm64", "aarch64"]
-
-
-# If Mac computer with Apple Silicon architecture, check if macOS version is 13.5 or above
-def is_mac_os_version_supported(min_version: str = "13.5") -> bool:
-    if not is_mac_environment() or not is_apple_silicon_cpu():
-        return False
-    mac_version = platform.mac_ver()[0]
-    if not mac_version:
-        return False
-    # print("Mac OS Version:", mac_version)
-    return version.parse(mac_version) >= version.parse(min_version)
-
-if __name__ == "__main__":
-    print("Is Mac Environment:", is_mac_environment())
-    print("Is Apple Silicon CPU:", is_apple_silicon_cpu())
-    print("Is Mac OS Version Supported (>=13.5):", is_mac_os_version_supported())

NaviOCR/tools/config_reader.py

@@ -1,136 +0,0 @@
-
-import json
-import os
-from loguru import logger
-
-try:
-    import torch
-    import torch_npu
-except ImportError:
-    pass
-
-
-# 定义配置文件名常量
-CONFIG_FILE_NAME = os.getenv('NaviOCR_TOOLS_CONFIG_JSON', 'NaviOCR.json')
-
-
-def read_config():
-    if os.path.isabs(CONFIG_FILE_NAME):
-        config_file = CONFIG_FILE_NAME
-    else:
-        home_dir = os.path.expanduser('~')
-        config_file = os.path.join(home_dir, CONFIG_FILE_NAME)
-
-    if not os.path.exists(config_file):
-        # logger.warning(f'{config_file} not found, using default configuration')
-        return None
-    else:
-        with open(config_file, 'r', encoding='utf-8') as f:
-            config = json.load(f)
-        return config
-
-
-def get_s3_config(bucket_name: str):
-    """~/magic-pdf.json 读出来."""
-    config = read_config()
-
-    bucket_info = config.get('bucket_info')
-    if bucket_name not in bucket_info:
-        access_key, secret_key, storage_endpoint = bucket_info['[default]']
-    else:
-        access_key, secret_key, storage_endpoint = bucket_info[bucket_name]
-
-    if access_key is None or secret_key is None or storage_endpoint is None:
-        raise Exception(f'ak, sk or endpoint not found in {CONFIG_FILE_NAME}')
-
-    # logger.info(f"get_s3_config: ak={access_key}, sk={secret_key}, endpoint={storage_endpoint}")
-
-    return access_key, secret_key, storage_endpoint
-
-
-def get_s3_config_dict(path: str):
-    access_key, secret_key, storage_endpoint = get_s3_config(get_bucket_name(path))
-    return {'ak': access_key, 'sk': secret_key, 'endpoint': storage_endpoint}
-
-
-def get_bucket_name(path):
-    bucket, key = parse_bucket_key(path)
-    return bucket
-
-
-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 get_device():
-    device_mode = os.getenv('NaviOCR_DEVICE_MODE', None)
-    if device_mode is not None:
-        return device_mode
-    else:
-        if torch.cuda.is_available():
-            return "cuda"
-        elif torch.backends.mps.is_available():
-            return "mps"
-        else:
-            try:
-                if torch_npu.npu.is_available():
-                    return "npu"
-            except Exception as e:
-                pass
-        return "cpu"
-
-
-def get_formula_enable(formula_enable):
-    formula_enable_env = os.getenv('NaviOCR_FORMULA_ENABLE')
-    formula_enable = formula_enable if formula_enable_env is None else formula_enable_env.lower() == 'true'
-    return formula_enable
-
-
-def get_table_enable(table_enable):
-    table_enable_env = os.getenv('NaviOCR_TABLE_ENABLE')
-    table_enable = table_enable if table_enable_env is None else table_enable_env.lower() == 'true'
-    return table_enable
-
-
-def get_latex_delimiter_config():
-    config = read_config()
-    if config is None:
-        return None
-    latex_delimiter_config = config.get('latex-delimiter-config', None)
-    if latex_delimiter_config is None:
-        # logger.warning(f"'latex-delimiter-config' not found in {CONFIG_FILE_NAME}, use 'None' as default")
-        return None
-    else:
-        return latex_delimiter_config
-
-
-def get_llm_aided_config():
-    config = read_config()
-    if config is None:
-        return None
-    llm_aided_config = config.get('llm-aided-config', None)
-    if llm_aided_config is None:
-        # logger.warning(f"'llm-aided-config' not found in {CONFIG_FILE_NAME}, use 'None' as default")
-        return None
-    else:
-        return llm_aided_config
-
-
-def get_local_models_dir():
-    config = read_config()
-    if config is None:
-        return None
-    models_dir = config.get('models-dir')
-    if models_dir is None:
-        logger.warning(f"'models-dir' not found in {CONFIG_FILE_NAME}, use None as default")
-    return models_dir

NaviOCR/tools/cut_image.py

@@ -1,36 +0,0 @@
-from loguru import logger
-
-from .pdf_image_tools import cut_image 
-import numpy as np
-
-def merge_bboxes_if_needed(bboxes):
-    if len(bboxes) > 4:
-        bboxes = np.array(bboxes).reshape(-1, 2)
-        minx = int(bboxes[:, 0].min())
-        miny = int(bboxes[:, 1].min())
-        maxx = int(bboxes[:, 0].max())
-        maxy = int(bboxes[:, 1].max())
-        return [minx, miny, maxx, maxy]
-    return [int(x) for x in bboxes]
-
-def cut_image_and_table(span, page_pil_img, page_img_md5, page_id, image_writer, scale=2):
-    def return_path(path_type):
-        return f"{path_type}/{page_img_md5}"
-
-    span_type = span["type"]
-    span["bbox"] = merge_bboxes_if_needed(span["bbox"])
-
-    if not check_img_bbox(span["bbox"]) or not image_writer:
-        span["image_path"] = ""
-    else:
-        span["image_path"] = cut_image(
-            span["bbox"], page_id, page_pil_img, return_path=return_path(span_type), image_writer=image_writer, scale=scale
-        )
-    return span
-
-
-def check_img_bbox(bbox) -> bool:
-    if any([bbox[0] >= bbox[2], bbox[1] >= bbox[3]]):
-        logger.warning(f"image_bboxes: 错误的box, {bbox}")
-        return False
-    return True

NaviOCR/tools/draw_bbox.py

@@ -1,634 +0,0 @@
-import json
-from io import BytesIO
-
-from loguru import logger
-from pypdf import PdfReader, PdfWriter, PageObject
-from reportlab.pdfgen import canvas
-
-from .enum_class import BlockType, ContentType, SplitFlag
-
-
-def cal_canvas_quad(page, bbox):
-    page_width, page_height = float(page.cropbox[2]), float(page.cropbox[3])
-    actual_width = page_width
-    actual_height = page_height
-    rotation_obj = page.get("/Rotate", 0)
-    try:
-        rotation = int(rotation_obj) % 360
-    except (ValueError, TypeError) as e:
-        logger.warning(f"Invalid /Rotate value {rotation_obj!r}; defaulting to 0. Error: {e}")
-        rotation = 0
-    if rotation in [90, 270]:
-        actual_width, actual_height = actual_height, actual_width
-
-    pts = []
-    for i in range(0, len(bbox), 2):
-        x = bbox[i]
-        y = bbox[i + 1]
-        if rotation == 270:
-            x, y = actual_height - y, actual_width - x
-        elif rotation == 180:
-            x = page_width - x
-        elif rotation == 90:
-            x, y = y, x
-        else:
-            y = page_height - y
-        pts.extend([x, y])
-    return pts
-
-
-def cal_canvas_rect(page, bbox):
-    """
-    Calculate the rectangle coordinates on the canvas based on the original PDF page and bounding box.
-
-    Args:
-        page: A PyPDF2 Page object representing a single page in the PDF.
-        bbox: [x0, y0, x1, y1] representing the bounding box coordinates.
-
-    Returns:
-        rect: [x0, y0, width, height] representing the rectangle coordinates on the canvas.
-    """
-    page_width, page_height = float(page.cropbox[2]), float(page.cropbox[3])
-    
-    actual_width = page_width    # The width of the final PDF display
-    actual_height = page_height  # The height of the final PDF display
-    
-    rotation_obj = page.get("/Rotate", 0)
-    try:
-        rotation = int(rotation_obj) % 360  # cast rotation to int to handle IndirectObject
-    except (ValueError, TypeError) as e:
-        logger.warning(f"Invalid /Rotate value {rotation_obj!r} on page; defaulting to 0. Error: {e}")
-        rotation = 0
-    
-    if rotation in [90, 270]:
-        # PDF is rotated 90 degrees or 270 degrees, and the width and height need to be swapped
-        actual_width, actual_height = actual_height, actual_width
-        
-    x0, y0, x1, y1 = bbox
-    rect_w = abs(x1 - x0)
-    rect_h = abs(y1 - y0)
-    
-    if rotation == 270:
-        rect_w, rect_h = rect_h, rect_w
-        x0 = actual_height - y1
-        y0 = actual_width - x1
-    elif rotation == 180:
-        x0 = page_width - x1
-        # y0 stays the same
-    elif rotation == 90:
-        rect_w, rect_h = rect_h, rect_w
-        x0, y0 = y0, x0 
-    else:
-        # rotation == 0
-        y0 = page_height - y1
-    
-    rect = [x0, y0, rect_w, rect_h]        
-    return rect
-
-
-def draw_bbox_without_number(i, bbox_list, page, c, rgb_config, fill_config):
-    new_rgb = [float(color) / 255 for color in rgb_config]
-    page_data = bbox_list[i]
-
-    for bbox in page_data:
-        # ---------- 2点框 ----------
-        if len(bbox) == 4:
-            rect = cal_canvas_rect(page, bbox)
-
-            if fill_config:
-                c.setFillColorRGB(new_rgb[0], new_rgb[1], new_rgb[2], 0.3)
-                c.rect(rect[0], rect[1], rect[2], rect[3], stroke=0, fill=1)
-            else:
-                c.setStrokeColorRGB(new_rgb[0], new_rgb[1], new_rgb[2])
-                c.rect(rect[0], rect[1], rect[2], rect[3], stroke=1, fill=0)
-        # ---------- 多点框 ----------
-        elif len(bbox) % 2 == 0:
-            pts = cal_canvas_quad(page, bbox)
-            path = c.beginPath()
-            path.moveTo(pts[0], pts[1])
-            for j in range(2, len(pts), 2):
-                path.lineTo(pts[j], pts[j + 1])
-            path.close()
-            if fill_config:
-                c.setFillColorRGB(*new_rgb, 0.3)
-                c.drawPath(path, stroke=0, fill=1)
-            else:
-                c.setStrokeColorRGB(*new_rgb)
-                c.drawPath(path, stroke=1, fill=0)
-        else:
-            logger.warning(f"Invalid bbox length: {bbox}")
-            continue
-
-    return c
-
-
-def draw_bbox_with_number(i, bbox_list, page, c, rgb_config, fill_config, draw_bbox=True):
-    new_rgb = [float(color) / 255 for color in rgb_config]
-    page_data = bbox_list[i]
-
-    page_width, page_height = float(page.cropbox[2]), float(page.cropbox[3])
-
-    for j, (bbox, angle) in enumerate(page_data):
-
-        # ---------- 2点框 ----------
-        if len(bbox) == 4:
-            rect = cal_canvas_rect(page, bbox)
-
-            if draw_bbox:
-                if fill_config:
-                    c.setFillColorRGB(*new_rgb, 0.3)
-                    c.rect(rect[0], rect[1], rect[2], rect[3], stroke=0, fill=1)
-                else:
-                    c.setStrokeColorRGB(*new_rgb)
-                    c.rect(rect[0], rect[1], rect[2], rect[3], stroke=1, fill=0)
-
-        # ---------- 多点框 ----------
-        elif len(bbox) % 2 == 0:
-            pts = cal_canvas_quad(page, bbox)
-
-            xs = pts[0::2]
-            ys = pts[1::2]
-
-            rect = [min(xs), min(ys), max(xs) - min(xs), max(ys) - min(ys)]
-
-            if draw_bbox:
-                path = c.beginPath()
-                path.moveTo(pts[0], pts[1])
-                for i in range(2, len(pts), 2):
-                    path.lineTo(pts[i], pts[i + 1])
-                path.close()
-                if fill_config:
-                    c.setFillColorRGB(*new_rgb, 0.3)
-                    c.drawPath(path, stroke=0, fill=1)
-                else:
-                    c.setStrokeColorRGB(*new_rgb)
-                    c.drawPath(path, stroke=1, fill=0)
-        else:
-            logger.warning(f"Invalid bbox length: {bbox}")
-            continue
-        c.setFillColorRGB(*new_rgb, 1.0)
-        c.setFontSize(size=10)
-
-        c.saveState()
-
-        rotation_obj = page.get("/Rotate", 0)
-        try:
-            rotation = int(rotation_obj) % 360
-        except (ValueError, TypeError):
-            logger.warning(f"Invalid /Rotate value: {rotation_obj!r}, defaulting to 0")
-            rotation = 0
-
-        if rotation == 0:
-            c.translate(rect[0] + rect[2] + 2, rect[1] + rect[3] - 10)
-        elif rotation == 90:
-            c.translate(rect[0] + 10, rect[1] + rect[3] + 2)
-        elif rotation == 180:
-            c.translate(rect[0] - 2, rect[1] + 10)
-        elif rotation == 270:
-            c.translate(rect[0] + rect[2] - 10, rect[1] - 2)
-
-        c.rotate(rotation)
-        c.drawString(0, 0, str(j + 1) + '+' + str(angle))
-
-        c.restoreState()
-
-    return c
-
-
-def draw_layout_bbox(pdf_info, pdf_bytes, out_path, filename):
-    dropped_bbox_list = []
-
-    tables_body_list, tables_caption_list, tables_footnote_list = [], [], []
-    imgs_body_list, imgs_caption_list, imgs_footnote_list = [], [], []
-    codes_body_list, codes_caption_list = [], []
-
-    titles_list = []
-    texts_list = []
-    interequations_list = []
-    lists_list = []
-    list_items_list = []
-    indexs_list = []
-    seal_list = []
-    char_list = []
-
-    for page in pdf_info:
-        page_dropped_list = []
-
-        tables_body, tables_caption, tables_footnote = [], [], []
-        imgs_body, imgs_caption, imgs_footnote = [], [], []
-        codes_body, codes_caption = [], []
-
-        titles = []
-        texts = []
-        interequations = []
-        lists = []
-        list_items = []
-        indices = []
-        seals = []
-        chars = []
-
-        for dropped_bbox in page['discarded_blocks']:
-            page_dropped_list.append(dropped_bbox['bbox'])
-
-        dropped_bbox_list.append(page_dropped_list)
-
-        for block in page["para_blocks"]:
-            block_bbox = block["bbox"]
-
-            if block["type"] == BlockType.TABLE:
-                for nested_block in block.get("blocks", []):
-                    nested_bbox = nested_block["bbox"]
-
-                    if nested_block["type"] == BlockType.TABLE_BODY:
-                        tables_body.append(nested_bbox)
-
-                    elif nested_block["type"] == BlockType.TABLE_CAPTION:
-                        tables_caption.append(nested_bbox)
-
-                    elif nested_block["type"] == BlockType.TABLE_FOOTNOTE:
-                        if nested_block.get(SplitFlag.CROSS_PAGE, False):
-                            continue
-                        tables_footnote.append(nested_bbox)
-
-            elif block["type"] == BlockType.IMAGE:
-                for nested_block in block.get("blocks", []):
-                    nested_bbox = nested_block["bbox"]
-
-                    if nested_block["type"] == BlockType.IMAGE_BODY:
-                        imgs_body.append(nested_bbox)
-
-                    elif nested_block["type"] == BlockType.IMAGE_CAPTION:
-                        imgs_caption.append(nested_bbox)
-
-                    elif nested_block["type"] == BlockType.IMAGE_FOOTNOTE:
-                        imgs_footnote.append(nested_bbox)
-
-            elif block["type"] == BlockType.CODE:
-                for nested_block in block.get("blocks", []):
-                    nested_bbox = nested_block["bbox"]
-
-                    if nested_block["type"] == BlockType.CODE_BODY:
-                        codes_body.append(nested_bbox)
-
-                    elif nested_block["type"] == BlockType.CODE_CAPTION:
-                        codes_caption.append(nested_bbox)
-
-            elif block["type"] == BlockType.TITLE:
-                titles.append(block_bbox)
-
-            elif block["type"] in [BlockType.TEXT, BlockType.REF_TEXT]:
-                texts.append(block_bbox)
-
-            elif block["type"] == BlockType.INTERLINE_EQUATION:
-                interequations.append(block_bbox)
-
-            elif block["type"] == BlockType.SEAL:
-                seals.append(block_bbox)
-
-            elif block["type"] == BlockType.CHAR:
-                chars.append(block_bbox)
-
-            elif block["type"] == BlockType.LIST:
-                lists.append(block_bbox)
-
-                if "blocks" in block:
-                    for sub_block in block["blocks"]:
-                        list_items.append(sub_block["bbox"])
-
-            elif block["type"] == BlockType.INDEX:
-                indices.append(block_bbox)
-
-        tables_body_list.append(tables_body)
-        tables_caption_list.append(tables_caption)
-        tables_footnote_list.append(tables_footnote)
-
-        imgs_body_list.append(imgs_body)
-        imgs_caption_list.append(imgs_caption)
-        imgs_footnote_list.append(imgs_footnote)
-
-        titles_list.append(titles)
-        texts_list.append(texts)
-        interequations_list.append(interequations)
-
-        seal_list.append(seals)
-        char_list.append(chars) 
-        lists_list.append(lists)
-        list_items_list.append(list_items)
-
-        indexs_list.append(indices)
-
-        codes_body_list.append(codes_body)
-        codes_caption_list.append(codes_caption)
-
-    layout_bbox_list = []
-
-    table_type_order = {
-        "table_caption": 1,
-        "table_body": 2,
-        "table_footnote": 3
-    }
-
-    for page in pdf_info:
-        page_block_list = []
-
-        for block in page["para_blocks"]:
-            if block["type"] in [
-                BlockType.TEXT,
-                BlockType.REF_TEXT,
-                BlockType.TITLE,
-                BlockType.INTERLINE_EQUATION,
-                BlockType.SEAL,
-                BlockType.CHAR,
-                BlockType.LIST,
-                BlockType.INDEX,
-            ]:
-                bbox = block["bbox"]
-                page_block_list.append([bbox, block["angle"]])
-
-            elif block["type"] == BlockType.IMAGE:
-                for sub_block in block.get("blocks", []):
-                    bbox = sub_block["bbox"]
-                    page_block_list.append([bbox, sub_block["angle"]])
-
-            elif block["type"] == BlockType.TABLE:
-                sorted_blocks = sorted(
-                    block.get("blocks", []),
-                    key=lambda x: table_type_order[x["type"]]
-                )
-
-                for sub_block in sorted_blocks:
-                    if sub_block.get(SplitFlag.CROSS_PAGE, False):
-                        continue
-
-                    bbox = sub_block["bbox"]
-                    page_block_list.append([bbox, sub_block["angle"]])
-
-            elif block["type"] == BlockType.CODE:
-                for sub_block in block.get("blocks", []):
-                    bbox = sub_block["bbox"]
-                    page_block_list.append([bbox, sub_block["angle"]])
-
-        layout_bbox_list.append(page_block_list)
-
-    pdf_bytes_io = BytesIO(pdf_bytes)
-    pdf_docs = PdfReader(pdf_bytes_io)
-    output_pdf = PdfWriter()
-
-    for i, page in enumerate(pdf_docs.pages):
-        page_width = float(page.cropbox[2])
-        page_height = float(page.cropbox[3])
-
-        custom_page_size = (page_width, page_height)
-
-        packet = BytesIO()
-        c = canvas.Canvas(packet, pagesize=custom_page_size)
-
-        c = draw_bbox_without_number(i, codes_body_list, page, c, [102, 0, 204], True)
-        c = draw_bbox_without_number(i, codes_caption_list, page, c, [204, 153, 255], True)
-        c = draw_bbox_without_number(i, dropped_bbox_list, page, c, [158, 158, 158], True)
-
-        c = draw_bbox_without_number(i, tables_body_list, page, c, [204, 204, 0], True)
-        c = draw_bbox_without_number(i, tables_caption_list, page, c, [255, 255, 102], True)
-        c = draw_bbox_without_number(i, tables_footnote_list, page, c, [229, 255, 204], True)
-
-        c = draw_bbox_without_number(i, imgs_body_list, page, c, [153, 255, 51], True)
-        c = draw_bbox_without_number(i, imgs_caption_list, page, c, [102, 178, 255], True)
-        c = draw_bbox_without_number(i, imgs_footnote_list, page, c, [255, 178, 102], True)
-
-        c = draw_bbox_without_number(i, titles_list, page, c, [102, 102, 255], True)
-        c = draw_bbox_without_number(i, texts_list, page, c, [153, 0, 76], True)
-
-        c = draw_bbox_without_number(i, interequations_list, page, c, [0, 255, 0], True)
-        c = draw_bbox_without_number(i, seal_list, page, c, [255, 0, 255], True)
-        c = draw_bbox_without_number(i, char_list, page, c, [255, 100, 255], True)
-
-        c = draw_bbox_without_number(i, lists_list, page, c, [40, 169, 92], True)
-        c = draw_bbox_without_number(i, list_items_list, page, c, [40, 169, 92], False)
-
-        c = draw_bbox_without_number(i, indexs_list, page, c, [40, 169, 92], True)
-
-        c = draw_bbox_with_number(
-            i,
-            layout_bbox_list,
-            page,
-            c,
-            [255, 0, 0],
-            False,
-            draw_bbox=False
-        )
-
-        c.save()
-
-        packet.seek(0)
-        overlay_pdf = PdfReader(packet)
-
-        if len(overlay_pdf.pages) > 0:
-            new_page = PageObject(pdf=None)
-            new_page.update(page)
-
-            page = new_page
-            page.merge_page(overlay_pdf.pages[0])
-
-        output_pdf.add_page(page)
-
-    with open(f"{out_path}/{filename}", "wb") as f:
-        output_pdf.write(f)
-
-def draw_span_bbox(pdf_info, pdf_bytes, out_path, filename):
-    text_list = []
-    inline_equation_list = []
-    interline_equation_list = []
-    image_list = []
-    table_list = []
-    dropped_list = []
-
-    def get_span_info(span):
-        if span['type'] == ContentType.TEXT:
-            page_text_list.append(span['bbox'])
-        elif span['type'] == ContentType.INLINE_EQUATION:
-            page_inline_equation_list.append(span['bbox'])
-        elif span['type'] == ContentType.INTERLINE_EQUATION:
-            page_interline_equation_list.append(span['bbox'])
-        elif span['type'] == ContentType.IMAGE:
-            page_image_list.append(span['bbox'])
-        elif span['type'] == ContentType.TABLE:
-            page_table_list.append(span['bbox'])
-
-    for page in pdf_info:
-        page_text_list = []
-        page_inline_equation_list = []
-        page_interline_equation_list = []
-        page_image_list = []
-        page_table_list = []
-        page_dropped_list = []
-
-
-        # 构造dropped_list
-        for block in page['discarded_blocks']:
-            if block['type'] == BlockType.DISCARDED:
-                for line in block['lines']:
-                    for span in line['spans']:
-                        page_dropped_list.append(span['bbox'])
-        dropped_list.append(page_dropped_list)
-        # 构造其余useful_list
-        for block in page['para_blocks']:  # span直接用分段合并前的结果就可以
-        # for block in page['preproc_blocks']:
-            # print(block.keys(), block['type'])
-            if block['type'] in [
-                BlockType.TEXT,
-                BlockType.TITLE,
-                BlockType.INTERLINE_EQUATION,
-                BlockType.INDEX,
-            ]:
-                for line in block['lines']:
-                    for span in line['spans']:
-                        get_span_info(span)
-            elif block['type'] in [BlockType.IMAGE, BlockType.TABLE, BlockType.LIST]:
-                for sub_block in block['blocks']:
-                    for line in sub_block['lines']:
-                        for span in line['spans']:
-                            get_span_info(span)
-        text_list.append(page_text_list)
-        inline_equation_list.append(page_inline_equation_list)
-        interline_equation_list.append(page_interline_equation_list)
-        image_list.append(page_image_list)
-        table_list.append(page_table_list)
-
-    pdf_bytes_io = BytesIO(pdf_bytes)
-    pdf_docs = PdfReader(pdf_bytes_io)
-    output_pdf = PdfWriter()
-
-    for i, page in enumerate(pdf_docs.pages):
-        # 获取原始页面尺寸
-        page_width, page_height = float(page.cropbox[2]), float(page.cropbox[3])
-        custom_page_size = (page_width, page_height)
-
-        packet = BytesIO()
-        # 使用原始PDF的尺寸创建canvas
-        c = canvas.Canvas(packet, pagesize=custom_page_size)
-
-        # 获取当前页面的数据
-        draw_bbox_without_number(i, text_list, page, c,[255, 0, 0], False)
-        draw_bbox_without_number(i, inline_equation_list, page, c, [0, 255, 0], False)
-        draw_bbox_without_number(i, interline_equation_list, page, c, [0, 0, 255], False)
-        draw_bbox_without_number(i, image_list, page, c, [255, 204, 0], False)
-        draw_bbox_without_number(i, table_list, page, c, [204, 0, 255], False)
-        draw_bbox_without_number(i, dropped_list, page, c, [158, 158, 158], False)
-
-        c.save()
-        packet.seek(0)
-        overlay_pdf = PdfReader(packet)
-
-        # 添加检查确保overlay_pdf.pages不为空
-        if len(overlay_pdf.pages) > 0:
-            new_page = PageObject(pdf=None)
-            new_page.update(page)
-            page = new_page
-            page.merge_page(overlay_pdf.pages[0])
-        else:
-            # 记录日志并继续处理下一个页面
-            # logger.warning(f"span.pdf: 第{i + 1}页未能生成有效的overlay PDF")
-            pass
-
-        output_pdf.add_page(page)
-
-    # Save the PDF
-    with open(f"{out_path}/{filename}", "wb") as f:
-        output_pdf.write(f)
-
-
-def draw_line_sort_bbox(pdf_info, pdf_bytes, out_path, filename):
-    layout_bbox_list = []
-
-    for page in pdf_info:
-        page_line_list = []
-        for block in page['preproc_blocks']:
-            if block['type'] in [BlockType.TEXT]:
-                for line in block['lines']:
-                    bbox = line['bbox']
-                    index = line['index']
-                    page_line_list.append({'index': index, 'bbox': bbox})
-            elif block['type'] in [BlockType.TITLE, BlockType.INTERLINE_EQUATION]:
-                if 'virtual_lines' in block:
-                    if len(block['virtual_lines']) > 0 and block['virtual_lines'][0].get('index', None) is not None:
-                        for line in block['virtual_lines']:
-                            bbox = line['bbox']
-                            index = line['index']
-                            page_line_list.append({'index': index, 'bbox': bbox})
-                else:
-                    for line in block['lines']:
-                        bbox = line['bbox']
-                        index = line['index']
-                        page_line_list.append({'index': index, 'bbox': bbox})
-            elif block['type'] in [BlockType.IMAGE, BlockType.TABLE]:
-                for sub_block in block['blocks']:
-                    if sub_block['type'] in [BlockType.IMAGE_BODY, BlockType.TABLE_BODY]:
-                        if len(sub_block['virtual_lines']) > 0 and sub_block['virtual_lines'][0].get('index', None) is not None:
-                            for line in sub_block['virtual_lines']:
-                                bbox = line['bbox']
-                                index = line['index']
-                                page_line_list.append({'index': index, 'bbox': bbox})
-                        else:
-                            for line in sub_block['lines']:
-                                bbox = line['bbox']
-                                index = line['index']
-                                page_line_list.append({'index': index, 'bbox': bbox})
-                    elif sub_block['type'] in [BlockType.IMAGE_CAPTION, BlockType.TABLE_CAPTION, BlockType.IMAGE_FOOTNOTE, BlockType.TABLE_FOOTNOTE]:
-                        for line in sub_block['lines']:
-                            bbox = line['bbox']
-                            index = line['index']
-                            page_line_list.append({'index': index, 'bbox': bbox})
-        sorted_bboxes = sorted(page_line_list, key=lambda x: x['index'])
-        layout_bbox_list.append(sorted_bbox['bbox'] for sorted_bbox in sorted_bboxes)
-    pdf_bytes_io = BytesIO(pdf_bytes)
-    pdf_docs = PdfReader(pdf_bytes_io)
-    output_pdf = PdfWriter()
-
-    for i, page in enumerate(pdf_docs.pages):
-        # 获取原始页面尺寸
-        page_width, page_height = float(page.cropbox[2]), float(page.cropbox[3])
-        custom_page_size = (page_width, page_height)
-
-        packet = BytesIO()
-        # 使用原始PDF的尺寸创建canvas
-        c = canvas.Canvas(packet, pagesize=custom_page_size)
-
-        # 获取当前页面的数据
-        draw_bbox_with_number(i, layout_bbox_list, page, c, [255, 0, 0], False)
-
-        c.save()
-        packet.seek(0)
-        overlay_pdf = PdfReader(packet)
-
-        # 添加检查确保overlay_pdf.pages不为空
-        if len(overlay_pdf.pages) > 0:
-            new_page = PageObject(pdf=None)
-            new_page.update(page)
-            page = new_page
-            page.merge_page(overlay_pdf.pages[0])
-        else:
-            # 记录日志并继续处理下一个页面
-            # logger.warning(f"span.pdf: 第{i + 1}页未能生成有效的overlay PDF")
-            pass
-
-        output_pdf.add_page(page)
-
-    # Save the PDF
-    with open(f"{out_path}/{filename}", "wb") as f:
-        output_pdf.write(f)
-
-
-if __name__ == "__main__":
-    # 读取PDF文件
-    pdf_path = "examples/demo1.pdf"
-    with open(pdf_path, "rb") as f:
-        pdf_bytes = f.read()
-
-    # 从json文件读取pdf_info
-
-    json_path = "examples/demo1_1746005777.0863056_middle.json"
-    with open(json_path, "r", encoding="utf-8") as f:
-        pdf_ann = json.load(f)
-    pdf_info = pdf_ann["pdf_info"]
-    # 调用可视化函数,输出到examples目录
-    draw_layout_bbox(pdf_info, pdf_bytes, "examples", "output_with_layout.pdf")

NaviOCR/tools/enum_class.py

@@ -1,67 +0,0 @@
-from enum import Enum
-
-class BlockType:
-    IMAGE = 'image'
-    TABLE = 'table'
-    IMAGE_BODY = 'image_body'
-    TABLE_BODY = 'table_body'
-    IMAGE_CAPTION = 'image_caption'
-    TABLE_CAPTION = 'table_caption'
-    IMAGE_FOOTNOTE = 'image_footnote'
-    TABLE_FOOTNOTE = 'table_footnote'
-    TEXT = 'text' 
-    TITLE = 'title'
-    INTERLINE_EQUATION = 'interline_equation'
-    LIST = 'list'
-    INDEX = 'index'
-    DISCARDED = 'discarded'
-
-    CODE = "code"
-    CODE_BODY = "code_body"
-    CODE_CAPTION = "code_caption"
-    ALGORITHM = "algorithm"
-    REF_TEXT = "ref_text"
-    PHONETIC = "phonetic"
-    HEADER = "header"
-    FOOTER = "footer"
-    PAGE_NUMBER = "page_number"
-    ASIDE_TEXT = "aside_text"
-    PAGE_FOOTNOTE = "page_footnote"
-    SEAL = 'seal'
-    CHAR = 'char'
-
-class ContentType:
-    IMAGE = 'image'
-    TABLE = 'table'
-    TEXT = 'text'
-    INTERLINE_EQUATION = 'interline_equation'
-    INLINE_EQUATION = 'inline_equation'
-    EQUATION = 'equation'
-    CODE = 'code'
-    SEAL = 'seal'
-    CHAR = 'char'
-
-class SplitFlag:
-    CROSS_PAGE = 'cross_page'
-    LINES_DELETED = 'lines_deleted'
-
-
-class ImageType:
-    PIL = 'pil_img'
-    BASE64 = 'base64_img'
-
-
-class NotExtractType(Enum):
-    TEXT = BlockType.TEXT
-    TITLE = BlockType.TITLE
-    HEADER = BlockType.HEADER
-    FOOTER = BlockType.FOOTER
-    PAGE_NUMBER = BlockType.PAGE_NUMBER
-    PAGE_FOOTNOTE = BlockType.PAGE_FOOTNOTE
-    REF_TEXT = BlockType.REF_TEXT
-    TABLE_CAPTION = BlockType.TABLE_CAPTION
-    IMAGE_CAPTION = BlockType.IMAGE_CAPTION
-    TABLE_FOOTNOTE = BlockType.TABLE_FOOTNOTE
-    IMAGE_FOOTNOTE = BlockType.IMAGE_FOOTNOTE
-    CODE_CAPTION = BlockType.CODE_CAPTION
-

NaviOCR/tools/guess_suffix_or_lang.py

@@ -1,55 +0,0 @@
-from pathlib import Path
-
-from loguru import logger
-from magika import Magika
-
-
-DEFAULT_LANG = "txt"
-PDF_SIG_BYTES = b'%PDF'
-magika = Magika()
-
-def code_content_clean(content):
-    if not content:
-        return ""
-    lines = content.splitlines()
-    start_idx = 0
-    end_idx = len(lines)
-    if lines and lines[0].startswith("```"):
-        start_idx = 1
-    if lines and end_idx > start_idx and lines[end_idx - 1].strip() == "```":
-        end_idx -= 1
-    if start_idx < end_idx:
-        return "\n".join(lines[start_idx:end_idx]).strip()
-    return ""
-
-def guess_language_by_text(code):
-    code = code_content_clean(code)
-    if code.startswith("<_"):
-        end = code.find("_>")
-        if end != -1:
-            lang = code[2:end].lower()
-            code = code[end + 2:]
-            return lang, code
-
-    return (lang if lang != "unknown" else DEFAULT_LANG), code
-
-
-def guess_suffix_by_bytes(file_bytes, file_path=None) -> str:
-    suffix = magika.identify_bytes(file_bytes).prediction.output.label
-    if file_path and suffix in ["ai", "html"] and Path(file_path).suffix.lower() in [".pdf"] and file_bytes[:4] == PDF_SIG_BYTES:
-        suffix = "pdf"
-    return suffix
-
-
-def guess_suffix_by_path(file_path) -> str:
-    if not isinstance(file_path, Path):
-        file_path = Path(file_path)
-    suffix = magika.identify_path(file_path).prediction.output.label
-    if suffix in ["ai", "html"] and file_path.suffix.lower() in [".pdf"]:
-        try:
-            with open(file_path, 'rb') as f:
-                if f.read(4) == PDF_SIG_BYTES:
-                    suffix = "pdf"
-        except Exception as e:
-            logger.warning(f"Failed to read file {file_path} for PDF signature check: {e}")
-    return suffix

NaviOCR/tools/hash_utils.py

@@ -1,30 +0,0 @@
-
-import hashlib
-import json
-
-
-def bytes_md5(file_bytes):
-    hasher = hashlib.md5()
-    hasher.update(file_bytes)
-    return hasher.hexdigest().upper()
-
-
-def str_md5(input_string):
-    hasher = hashlib.md5()
-    # 在Python3中，需要将字符串转化为字节对象才能被哈希函数处理
-    input_bytes = input_string.encode('utf-8')
-    hasher.update(input_bytes)
-    return hasher.hexdigest()
-
-
-def str_sha256(input_string):
-    hasher = hashlib.sha256()
-    # 在Python3中，需要将字符串转化为字节对象才能被哈希函数处理
-    input_bytes = input_string.encode('utf-8')
-    hasher.update(input_bytes)
-    return hasher.hexdigest()
-
-
-def dict_md5(d):
-    json_str = json.dumps(d, sort_keys=True, ensure_ascii=False)
-    return hashlib.md5(json_str.encode('utf-8')).hexdigest()

NaviOCR/tools/language.py

@@ -1,48 +0,0 @@
-import os
-import unicodedata
-
-if not os.getenv("FTLANG_CACHE"):
-    current_file_path = os.path.abspath(__file__)
-    current_dir = os.path.dirname(current_file_path)
-    root_dir = os.path.dirname(current_dir)
-    ftlang_cache_dir = os.path.join(root_dir, 'resources', 'fasttext-langdetect')
-    os.environ["FTLANG_CACHE"] = str(ftlang_cache_dir)
-    # print(os.getenv("FTLANG_CACHE"))
-
-from fast_langdetect import detect_language
-
-
-def remove_invalid_surrogates(text):
-    # 移除无效的 UTF-16 代理对
-    return ''.join(c for c in text if not (0xD800 <= ord(c) <= 0xDFFF))
-
-
-def detect_lang(text: str, MAX_LANG_LEN = 128) -> str:
-
-    if not text:
-        return ""
-
-    text = text.replace("\n", "")
-    text = remove_invalid_surrogates(text)
-
-    text = text[:MAX_LANG_LEN]
-    try:
-        lang_upper = detect_language(text)
-    except:
-        html_no_ctrl_chars = ''.join([l for l in text if unicodedata.category(l)[0] not in ['C', ]])
-        lang_upper = detect_language(html_no_ctrl_chars)
-
-    try:
-        lang = lang_upper.lower()
-    except:
-        lang = ""
-    return lang
-
-
-if __name__ == '__main__':
-    print(os.getenv("FTLANG_CACHE"))
-    print(detect_lang("This is a test."*100000))
-    print(detect_lang("<html>This is a test</html>"))
-    print(detect_lang("这个是中文测试。"))
-    print(detect_lang("<html>这个是中文测试。</html>"))
-    print(detect_lang("〖\ud835\udc46\ud835〗这是个包含utf-16的中文测试"))

NaviOCR/tools/magic_model_utils.py

@@ -1,251 +0,0 @@
-"""
-包含两个MagicModel类中重复使用的方法和逻辑
-"""
-from typing import List, Dict, Any, Callable
-from NaviOCR.tools.boxbase import bbox_distance, bbox_center_distance, is_in
-
-
-def reduct_overlap(bboxes: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-    """
-    去除重叠的bbox，保留不被其他bbox包含的bbox
-
-    Args:
-        bboxes: 包含bbox信息的字典列表
-
-    Returns:
-        去重后的bbox列表
-    """
-    N = len(bboxes)
-    keep = [True] * N
-    for i in range(N):
-        for j in range(N):
-            if i == j:
-                continue
-            if is_in(bboxes[i]['bbox'], bboxes[j]['bbox']):
-                keep[i] = False
-    return [bboxes[i] for i in range(N) if keep[i]]
-
-
-def tie_up_category_by_distance_v3(
-        get_subjects_func: Callable,
-        get_objects_func: Callable,
-        extract_subject_func: Callable = None,
-        extract_object_func: Callable = None
-):
-    """
-    通用的类别关联方法，用于将主体对象与客体对象进行关联
-
-    参数:
-        get_subjects_func: 函数，提取主体对象
-        get_objects_func: 函数，提取客体对象
-        extract_subject_func: 函数，自定义提取主体属性（默认使用bbox和其他属性）
-        extract_object_func: 函数，自定义提取客体属性（默认使用bbox和其他属性）
-
-    返回:
-        关联后的对象列表
-    """
-    subjects = get_subjects_func()
-    objects = get_objects_func()
-
-    # 如果没有提供自定义提取函数，使用默认函数
-    if extract_subject_func is None:
-        extract_subject_func = lambda x: x
-    if extract_object_func is None:
-        extract_object_func = lambda x: x
-
-    ret = []
-    N, M = len(subjects), len(objects)
-    subjects.sort(key=lambda x: x["bbox"][0] ** 2 + x["bbox"][1] ** 2)
-    objects.sort(key=lambda x: x["bbox"][0] ** 2 + x["bbox"][1] ** 2)
-
-    OBJ_IDX_OFFSET = 10000
-    SUB_BIT_KIND, OBJ_BIT_KIND = 0, 1
-
-    all_boxes_with_idx = [(i, SUB_BIT_KIND, sub["bbox"][0], sub["bbox"][1]) for i, sub in enumerate(subjects)] + [
-        (i + OBJ_IDX_OFFSET, OBJ_BIT_KIND, obj["bbox"][0], obj["bbox"][1]) for i, obj in enumerate(objects)
-    ]
-    seen_idx = set()
-    seen_sub_idx = set()
-
-    while N > len(seen_sub_idx):
-        candidates = []
-        for idx, kind, x0, y0 in all_boxes_with_idx:
-            if idx in seen_idx:
-                continue
-            candidates.append((idx, kind, x0, y0))
-
-        if len(candidates) == 0:
-            break
-        left_x = min([v[2] for v in candidates])
-        top_y = min([v[3] for v in candidates])
-
-        candidates.sort(key=lambda x: (x[2] - left_x) ** 2 + (x[3] - top_y) ** 2)
-
-        fst_idx, fst_kind, left_x, top_y = candidates[0]
-        fst_bbox = subjects[fst_idx]['bbox'] if fst_kind == SUB_BIT_KIND else objects[fst_idx - OBJ_IDX_OFFSET]['bbox']
-        candidates.sort(
-            key=lambda x: bbox_distance(fst_bbox, subjects[x[0]]['bbox']) if x[1] == SUB_BIT_KIND else bbox_distance(
-                fst_bbox, objects[x[0] - OBJ_IDX_OFFSET]['bbox']))
-        nxt = None
-
-        for i in range(1, len(candidates)):
-            if candidates[i][1] ^ fst_kind == 1:
-                nxt = candidates[i]
-                break
-        if nxt is None:
-            break
-
-        if fst_kind == SUB_BIT_KIND:
-            sub_idx, obj_idx = fst_idx, nxt[0] - OBJ_IDX_OFFSET
-        else:
-            sub_idx, obj_idx = nxt[0], fst_idx - OBJ_IDX_OFFSET
-
-        pair_dis = bbox_distance(subjects[sub_idx]["bbox"], objects[obj_idx]["bbox"])
-        nearest_dis = float("inf")
-        for i in range(N):
-            # 取消原先算法中 1对1 匹配的偏置
-            # if i in seen_idx or i == sub_idx:continue
-            nearest_dis = min(nearest_dis, bbox_distance(subjects[i]["bbox"], objects[obj_idx]["bbox"]))
-
-        if pair_dis >= 3 * nearest_dis:
-            seen_idx.add(sub_idx)
-            continue
-
-        seen_idx.add(sub_idx)
-        seen_idx.add(obj_idx + OBJ_IDX_OFFSET)
-        seen_sub_idx.add(sub_idx)
-
-        ret.append(
-            {
-                "sub_bbox": extract_subject_func(subjects[sub_idx]),
-                "obj_bboxes": [extract_object_func(objects[obj_idx])],
-                "sub_idx": sub_idx,
-            }
-        )
-
-    for i in range(len(objects)):
-        j = i + OBJ_IDX_OFFSET
-        if j in seen_idx:
-            continue
-        seen_idx.add(j)
-        nearest_dis, nearest_sub_idx = float("inf"), -1
-        for k in range(len(subjects)):
-            dis = bbox_distance(objects[i]["bbox"], subjects[k]["bbox"])
-            if dis < nearest_dis:
-                nearest_dis = dis
-                nearest_sub_idx = k
-
-        for k in range(len(subjects)):
-            if k != nearest_sub_idx:
-                continue
-            if k in seen_sub_idx:
-                for kk in range(len(ret)):
-                    if ret[kk]["sub_idx"] == k:
-                        ret[kk]["obj_bboxes"].append(extract_object_func(objects[i]))
-                        break
-            else:
-                ret.append(
-                    {
-                        "sub_bbox": extract_subject_func(subjects[k]),
-                        "obj_bboxes": [extract_object_func(objects[i])],
-                        "sub_idx": k,
-                    }
-                )
-            seen_sub_idx.add(k)
-            seen_idx.add(k)
-
-    for i in range(len(subjects)):
-        if i in seen_sub_idx:
-            continue
-        ret.append(
-            {
-                "sub_bbox": extract_subject_func(subjects[i]),
-                "obj_bboxes": [],
-                "sub_idx": i,
-            }
-        )
-
-    return ret
-
-
-def tie_up_category_by_index(
-        get_subjects_func: Callable,
-        get_objects_func: Callable,
-        extract_subject_func: Callable = None,
-        extract_object_func: Callable = None
-):
-    """
-    基于index的类别关联方法，用于将主体对象与客体对象进行关联
-    客体优先匹配给index最接近的主体，index差值相同时使用bbox中心点距离作为tiebreaker
-
-    参数:
-        get_subjects_func: 函数，提取主体对象
-        get_objects_func: 函数，提取客体对象
-        extract_subject_func: 函数，自定义提取主体属性（默认使用bbox和其他属性）
-        extract_object_func: 函数，自定义提取客体属性（默认使用bbox和其他属性）
-
-    返回:
-        关联后的对象列表，按主体index升序排列
-    """
-    subjects = get_subjects_func()
-    objects = get_objects_func()
-
-    # 如果没有提供自定义提取函数，使用默认函数
-    if extract_subject_func is None:
-        extract_subject_func = lambda x: x
-    if extract_object_func is None:
-        extract_object_func = lambda x: x
-
-    # 初始化结果字典，key为主体索引，value为关联信息
-    result_dict = {}
-
-    # 初始化所有主体
-    for i, subject in enumerate(subjects):
-        result_dict[i] = {
-            "sub_bbox": extract_subject_func(subject),
-            "obj_bboxes": [],
-            "sub_idx": i,
-        }
-
-    # 为每个客体找到最匹配的主体
-    for obj in objects:
-        if len(subjects) == 0:
-            # 如果没有主体，跳过客体
-            continue
-
-        obj_index = obj["index"]
-        min_index_diff = float("inf")
-        best_subject_indices = []
-
-        # 找出index差值最小的所有主体
-        for i, subject in enumerate(subjects):
-            sub_index = subject["index"]
-            index_diff = abs(obj_index - sub_index)
-
-            if index_diff < min_index_diff:
-                min_index_diff = index_diff
-                best_subject_indices = [i]
-            elif index_diff == min_index_diff:
-                best_subject_indices.append(i)
-
-        # 如果有多个主体的index差值相同，使用中心点距离作为tiebreaker
-        if len(best_subject_indices) > 1:
-            min_center_dist = float("inf")
-            best_subject_idx = best_subject_indices[0]
-
-            for idx in best_subject_indices:
-                center_dist = bbox_center_distance(obj["bbox"], subjects[idx]["bbox"])
-                if center_dist < min_center_dist:
-                    min_center_dist = center_dist
-                    best_subject_idx = idx
-        else:
-            best_subject_idx = best_subject_indices[0]
-
-        # 将客体添加到最佳主体的obj_bboxes中
-        result_dict[best_subject_idx]["obj_bboxes"].append(extract_object_func(obj))
-
-    # 转换为列表并按主体index排序
-    ret = list(result_dict.values())
-    ret.sort(key=lambda x: x["sub_idx"])
-
-    return ret

NaviOCR/tools/model_utils.py

@@ -1,462 +0,0 @@
-import os
-import time
-import gc
-from PIL import Image
-from loguru import logger
-import numpy as np
-
-from NaviOCR.tools.boxbase import get_minbox_if_overlap_by_ratio
-
-try:
-    import torch
-    import torch_npu
-except ImportError:
-    pass
-
-
-def crop_img(input_res, input_img, crop_paste_x=0, crop_paste_y=0):
-
-    crop_xmin, crop_ymin = int(input_res['poly'][0]), int(input_res['poly'][1])
-    crop_xmax, crop_ymax = int(input_res['poly'][4]), int(input_res['poly'][5])
-
-    # Calculate new dimensions
-    crop_new_width = crop_xmax - crop_xmin + crop_paste_x * 2
-    crop_new_height = crop_ymax - crop_ymin + crop_paste_y * 2
-
-    if isinstance(input_img, np.ndarray):
-
-        # Create a white background array
-        return_image = np.ones((crop_new_height, crop_new_width, 3), dtype=np.uint8) * 255
-
-        # Crop the original image using numpy slicing
-        cropped_img = input_img[crop_ymin:crop_ymax, crop_xmin:crop_xmax]
-
-        # Paste the cropped image onto the white background
-        return_image[crop_paste_y:crop_paste_y + (crop_ymax - crop_ymin),
-        crop_paste_x:crop_paste_x + (crop_xmax - crop_xmin)] = cropped_img
-    else:
-        # Create a white background array
-        return_image = Image.new('RGB', (crop_new_width, crop_new_height), 'white')
-        # Crop image
-        crop_box = (crop_xmin, crop_ymin, crop_xmax, crop_ymax)
-        cropped_img = input_img.crop(crop_box)
-        return_image.paste(cropped_img, (crop_paste_x, crop_paste_y))
-
-    return_list = [crop_paste_x, crop_paste_y, crop_xmin, crop_ymin, crop_xmax, crop_ymax, crop_new_width,
-                   crop_new_height]
-    return return_image, return_list
-
-
-def get_coords_and_area(block_with_poly):
-    """Extract coordinates and area from a table."""
-    xmin, ymin = int(block_with_poly['poly'][0]), int(block_with_poly['poly'][1])
-    xmax, ymax = int(block_with_poly['poly'][4]), int(block_with_poly['poly'][5])
-    area = (xmax - xmin) * (ymax - ymin)
-    return xmin, ymin, xmax, ymax, area
-
-
-def calculate_intersection(box1, box2):
-    """Calculate intersection coordinates between two boxes."""
-    intersection_xmin = max(box1[0], box2[0])
-    intersection_ymin = max(box1[1], box2[1])
-    intersection_xmax = min(box1[2], box2[2])
-    intersection_ymax = min(box1[3], box2[3])
-
-    # Check if intersection is valid
-    if intersection_xmax <= intersection_xmin or intersection_ymax <= intersection_ymin:
-        return None
-
-    return intersection_xmin, intersection_ymin, intersection_xmax, intersection_ymax
-
-
-def calculate_iou(box1, box2):
-    """Calculate IoU between two boxes."""
-    intersection = calculate_intersection(box1[:4], box2[:4])
-
-    if not intersection:
-        return 0
-
-    intersection_xmin, intersection_ymin, intersection_xmax, intersection_ymax = intersection
-    intersection_area = (intersection_xmax - intersection_xmin) * (intersection_ymax - intersection_ymin)
-
-    area1, area2 = box1[4], box2[4]
-    union_area = area1 + area2 - intersection_area
-
-    return intersection_area / union_area if union_area > 0 else 0
-
-
-def is_inside(small_box, big_box, overlap_threshold=0.8):
-    """Check if small_box is inside big_box by at least overlap_threshold."""
-    intersection = calculate_intersection(small_box[:4], big_box[:4])
-
-    if not intersection:
-        return False
-
-    intersection_xmin, intersection_ymin, intersection_xmax, intersection_ymax = intersection
-    intersection_area = (intersection_xmax - intersection_xmin) * (intersection_ymax - intersection_ymin)
-
-    # Check if overlap exceeds threshold
-    return intersection_area >= overlap_threshold * small_box[4]
-
-
-def do_overlap(box1, box2):
-    """Check if two boxes overlap."""
-    return calculate_intersection(box1[:4], box2[:4]) is not None
-
-
-def merge_high_iou_tables(table_res_list, layout_res, table_indices, iou_threshold=0.7):
-    """Merge tables with IoU > threshold."""
-    if len(table_res_list) < 2:
-        return table_res_list, table_indices
-
-    table_info = [get_coords_and_area(table) for table in table_res_list]
-    merged = True
-
-    while merged:
-        merged = False
-        i = 0
-        while i < len(table_res_list) - 1:
-            j = i + 1
-            while j < len(table_res_list):
-                iou = calculate_iou(table_info[i], table_info[j])
-
-                if iou > iou_threshold:
-                    # Merge tables by taking their union
-                    x1_min, y1_min, x1_max, y1_max, _ = table_info[i]
-                    x2_min, y2_min, x2_max, y2_max, _ = table_info[j]
-
-                    union_xmin = min(x1_min, x2_min)
-                    union_ymin = min(y1_min, y2_min)
-                    union_xmax = max(x1_max, x2_max)
-                    union_ymax = max(y1_max, y2_max)
-
-                    # Create merged table
-                    merged_table = table_res_list[i].copy()
-                    merged_table['poly'] = [
-                        union_xmin, union_ymin, union_xmax, union_ymin,
-                        union_xmax, union_ymax, union_xmin, union_ymax
-                    ]
-                    # Update layout_res
-                    to_remove = [table_indices[j], table_indices[i]]
-                    for idx in sorted(to_remove, reverse=True):
-                        del layout_res[idx]
-                    layout_res.append(merged_table)
-
-                    # Update tracking lists
-                    table_indices = [k if k < min(to_remove) else
-                                     k - 1 if k < max(to_remove) else
-                                     k - 2 if k > max(to_remove) else
-                                     len(layout_res) - 1
-                                     for k in table_indices
-                                     if k not in to_remove]
-                    table_indices.append(len(layout_res) - 1)
-
-                    # Update table lists
-                    table_res_list.pop(j)
-                    table_res_list.pop(i)
-                    table_res_list.append(merged_table)
-
-                    # Update table_info
-                    table_info = [get_coords_and_area(table) for table in table_res_list]
-
-                    merged = True
-                    break
-                j += 1
-
-            if merged:
-                break
-            i += 1
-
-    return table_res_list, table_indices
-
-
-def filter_nested_tables(table_res_list, overlap_threshold=0.8, area_threshold=0.8):
-    """Remove big tables containing multiple smaller tables within them."""
-    if len(table_res_list) < 3:
-        return table_res_list
-
-    table_info = [get_coords_and_area(table) for table in table_res_list]
-    big_tables_idx = []
-
-    for i in range(len(table_res_list)):
-        # Find tables inside this one
-        tables_inside = [j for j in range(len(table_res_list))
-                         if i != j and is_inside(table_info[j], table_info[i], overlap_threshold)]
-
-        # Continue if there are at least 3 tables inside
-        if len(tables_inside) >= 3:
-            # Check if inside tables overlap with each other
-            tables_overlap = any(do_overlap(table_info[tables_inside[idx1]], table_info[tables_inside[idx2]])
-                                 for idx1 in range(len(tables_inside))
-                                 for idx2 in range(idx1 + 1, len(tables_inside)))
-
-            # If no overlaps, check area condition
-            if not tables_overlap:
-                total_inside_area = sum(table_info[j][4] for j in tables_inside)
-                big_table_area = table_info[i][4]
-
-                if total_inside_area > area_threshold * big_table_area:
-                    big_tables_idx.append(i)
-
-    return [table for i, table in enumerate(table_res_list) if i not in big_tables_idx]
-
-
-def remove_overlaps_min_blocks(res_list):
-
-    for res in res_list:
-        res['bbox'] = [int(res['poly'][0]), int(res['poly'][1]), int(res['poly'][4]), int(res['poly'][5])]
-
-    # 重叠block，小的不能直接删除，需要和大的那个合并成一个更大的。
-    # 删除重叠blocks中较小的那些
-    need_remove = []
-    for i in range(len(res_list)):
-        # 如果当前元素已在需要移除列表中，则跳过
-        if res_list[i] in need_remove:
-            continue
-
-        for j in range(i + 1, len(res_list)):
-            # 如果比较对象已在需要移除列表中，则跳过
-            if res_list[j] in need_remove:
-                continue
-
-            overlap_box = get_minbox_if_overlap_by_ratio(
-                res_list[i]['bbox'], res_list[j]['bbox'], 0.8
-            )
-
-            if overlap_box is not None:
-
-                # 根据重叠框确定哪个是小块，哪个是大块
-                if overlap_box == res_list[i]['bbox']:
-                    small_res, large_res = res_list[i], res_list[j]
-                elif overlap_box == res_list[j]['bbox']:
-                    small_res, large_res = res_list[j], res_list[i]
-                else:
-                    continue  # 如果重叠框与任一块都不匹配，跳过处理
-
-                if small_res['score'] <= large_res['score']:
-                    # 如果小块的分数低于大块，则小块为需要移除的块
-                    if small_res is not None and small_res not in need_remove:
-                        # 更新大块的边界为两者的并集
-                        x1, y1, x2, y2 = large_res['bbox']
-                        sx1, sy1, sx2, sy2 = small_res['bbox']
-                        x1 = min(x1, sx1)
-                        y1 = min(y1, sy1)
-                        x2 = max(x2, sx2)
-                        y2 = max(y2, sy2)
-                        large_res['bbox'] = [x1, y1, x2, y2]
-                        need_remove.append(small_res)
-                else:
-                    # 如果大块的分数低于小块，则大块为需要移除的块, 这时不需要更新小块的边界
-                    if large_res is not None and large_res not in need_remove:
-                        need_remove.append(large_res)
-
-    # 从列表中移除标记的元素
-    for res in need_remove:
-        res_list.remove(res)
-        del res['bbox']  # 删除bbox字段
-
-    for res in res_list:
-        # 将res的poly使用bbox重构
-        res['poly'] = [res['bbox'][0], res['bbox'][1], res['bbox'][2], res['bbox'][1],
-                       res['bbox'][2], res['bbox'][3], res['bbox'][0], res['bbox'][3]]
-        # 删除res的bbox
-        del res['bbox']
-
-    return res_list, need_remove
-
-
-def remove_overlaps_low_confidence_blocks(combined_res_list, overlap_threshold=0.8):
-    """
-    Remove low-confidence blocks that overlap with other blocks.
-
-    This function identifies and removes blocks with low confidence scores that overlap
-    with other blocks. It calculates the coordinates and area of each block, and checks
-    for overlaps based on a specified threshold. Blocks that meet the criteria for removal
-    are returned in a list.
-
-    Parameters:
-        combined_res_list (list): A list of blocks, where each block is a dictionary containing
-            keys like 'poly' (polygon coordinates) and optionally 'score' (confidence score).
-        overlap_threshold (float): The threshold for determining overlap between blocks. Default is 0.8.
-
-    Returns:
-        list: A list of blocks to be removed, based on the overlap and confidence criteria.
-    """
-    # 计算每个block的坐标和面积
-    block_info = []
-    for block in combined_res_list:
-        xmin, ymin = int(block['poly'][0]), int(block['poly'][1])
-        xmax, ymax = int(block['poly'][4]), int(block['poly'][5])
-        area = (xmax - xmin) * (ymax - ymin)
-        score = block.get('score', 0.5)  # 如果没有score字段，默认为0.5
-        block_info.append((xmin, ymin, xmax, ymax, area, score, block))
-
-    blocks_to_remove = []
-    marked_indices = set()  # 跟踪已标记为删除的block索引
-
-    # 检查每个block内部是否有3个及以上的小block
-    for i, (xmin, ymin, xmax, ymax, area, score, block) in enumerate(block_info):
-        # 如果当前block已标记为删除，则跳过
-        if i in marked_indices:
-            continue
-
-        # 查找内部的小block (仅考虑尚未被标记为删除的block)
-        blocks_inside = [(j, j_score, j_block) for j, (xj_min, yj_min, xj_max, yj_max, j_area, j_score, j_block) in
-                         enumerate(block_info)
-                         if i != j and j not in marked_indices and is_inside(block_info[j], block_info[i],
-                                                                             overlap_threshold)]
-
-        # 如果内部有3个及以上的小block
-        if len(blocks_inside) >= 2:
-            # 计算小block的平均分数
-            avg_score = sum(s for _, s, _ in blocks_inside) / len(blocks_inside)
-
-            # 比较大block的分数和小block的平均分数
-            if score > avg_score:
-                # 保留大block，扩展其边界
-                # 首先将所有小block标记为要删除
-                for j, _, j_block in blocks_inside:
-                    if j_block not in blocks_to_remove:
-                        blocks_to_remove.append(j_block)
-                        marked_indices.add(j)  # 标记索引为已处理
-
-                # 扩展大block的边界以包含所有小block
-                new_xmin, new_ymin, new_xmax, new_ymax = xmin, ymin, xmax, ymax
-                for _, _, j_block in blocks_inside:
-                    j_xmin, j_ymin = int(j_block['poly'][0]), int(j_block['poly'][1])
-                    j_xmax, j_ymax = int(j_block['poly'][4]), int(j_block['poly'][5])
-                    new_xmin = min(new_xmin, j_xmin)
-                    new_ymin = min(new_ymin, j_ymin)
-                    new_xmax = max(new_xmax, j_xmax)
-                    new_ymax = max(new_ymax, j_ymax)
-
-                # 更新大block的边界
-                block['poly'][0] = block['poly'][6] = new_xmin
-                block['poly'][1] = block['poly'][3] = new_ymin
-                block['poly'][2] = block['poly'][4] = new_xmax
-                block['poly'][5] = block['poly'][7] = new_ymax
-            else:
-                # 保留小blocks，删除大block
-                blocks_to_remove.append(block)
-                marked_indices.add(i)  # 标记当前索引为已处理
-    return blocks_to_remove
-
-
-def get_res_list_from_layout_res(layout_res, iou_threshold=0.7, overlap_threshold=0.8, area_threshold=0.8):
-    """Extract OCR, table and other regions from layout results."""
-    ocr_res_list = []
-    text_res_list = []
-    table_res_list = []
-    table_indices = []
-    single_page_mfdetrec_res = []
-
-    # Categorize regions
-    for i, res in enumerate(layout_res):
-        category_id = int(res['category_id'])
-
-        if category_id in [13, 14]:  # Formula regions
-            single_page_mfdetrec_res.append({
-                "bbox": [int(res['poly'][0]), int(res['poly'][1]),
-                         int(res['poly'][4]), int(res['poly'][5])],
-            })
-        elif category_id in [0, 2, 4, 6, 7, 3]:  # OCR regions
-            ocr_res_list.append(res)
-        elif category_id == 5:  # Table regions
-            table_res_list.append(res)
-            table_indices.append(i)
-        elif category_id in [1]:  # Text regions
-            text_res_list.append(res)
-
-    # Process tables: merge high IoU tables first, then filter nested tables
-    table_res_list, table_indices = merge_high_iou_tables(
-        table_res_list, layout_res, table_indices, iou_threshold)
-
-    filtered_table_res_list = filter_nested_tables(
-        table_res_list, overlap_threshold, area_threshold)
-
-    filtered_table_res_list, table_need_remove = remove_overlaps_min_blocks(filtered_table_res_list)
-
-    for res in table_need_remove:
-        if res in layout_res:
-            layout_res.remove(res)
-
-    # Remove filtered out tables from layout_res
-    if len(filtered_table_res_list) < len(table_res_list):
-        kept_tables = set(id(table) for table in filtered_table_res_list)
-        tables_to_remove = [table for table in table_res_list if id(table) not in kept_tables]
-        for table in tables_to_remove:
-            if table in layout_res:
-                layout_res.remove(table)
-
-    # Remove overlaps in OCR and text regions
-    text_res_list, need_remove = remove_overlaps_min_blocks(text_res_list)
-
-    ocr_res_list.extend(text_res_list)
-
-    for res in need_remove:
-        if res in layout_res:
-            layout_res.remove(res)
-
-    # 检测大block内部是否包含多个小block, 合并ocr和table列表进行检测
-    combined_res_list = ocr_res_list + filtered_table_res_list
-    blocks_to_remove = remove_overlaps_low_confidence_blocks(combined_res_list, overlap_threshold)
-    # 移除需要删除的blocks
-    for block in blocks_to_remove:
-        if block in ocr_res_list:
-            ocr_res_list.remove(block)
-        elif block in filtered_table_res_list:
-            filtered_table_res_list.remove(block)
-        # 同时从layout_res中删除
-        if block in layout_res:
-            layout_res.remove(block)
-
-    return ocr_res_list, filtered_table_res_list, single_page_mfdetrec_res
-
-
-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"):
-        if torch_npu.npu.is_available():
-            torch_npu.npu.empty_cache()
-    elif str(device).startswith("mps"):
-        torch.mps.empty_cache()
-    gc.collect()
-
-
-def clean_vram(device, vram_threshold=8):
-    total_memory = get_vram(device)
-    if total_memory and total_memory <= vram_threshold:
-        gc_start = time.time()
-        clean_memory(device)
-        gc_time = round(time.time() - gc_start, 2)
-        # logger.info(f"gc time: {gc_time}")
-
-
-def get_vram(device) -> int:
-    env_vram = os.getenv("NaviOCR_VIRTUAL_VRAM_SIZE")
-
-    # 如果环境变量已配置,尝试解析并返回
-    if env_vram is not None:
-        try:
-            total_memory = int(env_vram)
-            if total_memory > 0:
-                return total_memory
-            else:
-                logger.warning(
-                    f"NaviOCR_VIRTUAL_VRAM_SIZE value '{env_vram}' is not positive, falling back to auto-detection")
-        except ValueError:
-            logger.warning(
-                f"NaviOCR_VIRTUAL_VRAM_SIZE value '{env_vram}' is not a valid integer, falling back to auto-detection")
-
-    # 环境变量未配置或配置错误,根据device自动获取
-    total_memory = 1
-    if torch.cuda.is_available() and str(device).startswith("cuda"):
-        total_memory = round(torch.cuda.get_device_properties(device).total_memory / (1024 ** 3))  # 将字节转换为 GB
-    elif str(device).startswith("npu"):
-        if torch_npu.npu.is_available():
-            total_memory = round(torch_npu.npu.get_device_properties(device).total_memory / (1024 ** 3))  # 转为 GB
-
-    return total_memory

NaviOCR/tools/os_env_config.py

@@ -1,30 +0,0 @@
-import os
-
-
-def get_op_num_threads(env_name: str) -> int:
-    env_value = os.getenv(env_name, None)
-    return get_value_from_string(env_value, -1)
-
-
-def get_load_images_timeout() -> int:
-    env_value = os.getenv('NaviOCR_PDF_RENDER_TIMEOUT', None)
-    return get_value_from_string(env_value, 300)
-
-
-def get_value_from_string(env_value: str, default_value: int) -> int:
-    if env_value is not None:
-        try:
-            num_threads = int(env_value)
-            if num_threads > 0:
-                return num_threads
-        except ValueError:
-            return default_value
-    return default_value
-
-
-if __name__ == '__main__':
-    print(get_value_from_string('1', -1))
-    print(get_value_from_string('0', -1))
-    print(get_value_from_string('-1', -1))
-    print(get_value_from_string('abc', -1))
-    print(get_load_images_timeout())

NaviOCR/tools/pdf_image_tools.py

@@ -1,26 +0,0 @@
-import importlib
-
-import NaviOCR.config as CONFIG
-
-PDF_MODULES = {
-    "PyMuPDF": ".pdf_image_tools_PyMuPDF",
-    "pypdfium2": ".pdf_image_tools_pdfium",
-}
-
-module = importlib.import_module(
-    PDF_MODULES[CONFIG.PDF_TOOLS],
-    package=__package__,
-)
-
-for name in (
-    "pdf_page_to_image",
-    "_load_images_from_pdf_worker",
-    "load_images_from_pdf",
-    "load_images_from_pdf_core",
-    "cut_image",
-    "get_crop_img",
-    "images_bytes_to_pdf_bytes",
-    "get_page_size",
-    "convert_pdf_bytes_to_bytes"
-):
-    globals()[name] = getattr(module, name)

NaviOCR/tools/pdf_image_tools_PyMuPDF.py

@@ -1,289 +0,0 @@
-
-import os
-from io import BytesIO
-from concurrent.futures import ProcessPoolExecutor, TimeoutError as FuturesTimeoutError
-
-import fitz  # PyMuPDF
-import numpy as np
-from loguru import logger
-from PIL import Image, ImageOps
-
-import NaviOCR.config as CONFIG
-from NaviOCR.tools.check_sys_env import is_windows_environment
-from NaviOCR.tools.enum_class import ImageType
-from NaviOCR.tools.hash_utils import str_sha256
-from NaviOCR.tools.os_env_config import get_load_images_timeout
-from NaviOCR.tools.pdf_page_id import get_end_page_id
-from NaviOCR.tools.pdf_reader import image_to_b64str, image_to_bytes
-
-def convert_pdf_bytes_to_bytes(pdf_bytes,valid_single_page_ids=None):
-    pdf = fitz.open(stream=pdf_bytes, filetype="pdf")
-    output_pdf = fitz.open()
-    try:
-        total_pages = len(pdf)
-
-        if not valid_single_page_ids:
-            valid_single_page_ids = list(range(total_pages))
-
-        valid_single_page_ids = [
-            page_id
-            for page_id in valid_single_page_ids
-            if 0 <= page_id < total_pages
-        ]
-
-        for page_index in valid_single_page_ids:
-            try:
-                output_pdf.insert_pdf(
-                    pdf,
-                    from_page=page_index,
-                    to_page=page_index,
-                )
-            except Exception as page_error:
-                logger.warning(
-                    f"Failed to import page {page_index}: "
-                    f"{page_error}, skipping this page."
-                )
-
-        output_bytes = output_pdf.tobytes(
-            garbage=4,
-            deflate=True,
-        )
-
-    except Exception as e:
-        logger.warning(
-            f"Error in converting PDF bytes: {e}, "
-            f"Using original PDF bytes."
-        )
-        output_bytes = pdf_bytes
-
-    finally:
-        pdf.close()
-        output_pdf.close()
-
-    return output_bytes
-
-def pdf_page_to_image(page: fitz.Page,dpi=200,image_type=ImageType.PIL) -> dict:
-    zoom = dpi / 72.0
-    matrix = fitz.Matrix(zoom, zoom)
-    pix = page.get_pixmap(matrix=matrix, alpha=False)
-    pil_img = Image.frombytes("RGB",(pix.width, pix.height),pix.samples,)
-
-    image_dict = {"scale": zoom,}
-
-    if image_type == ImageType.BASE64:
-        image_dict["img_base64"] = image_to_b64str(pil_img)
-    else:
-        image_dict["img_pil"] = pil_img
-
-    return image_dict
-
-
-def _load_images_from_pdf_worker(
-    pdf_bytes,
-    dpi,
-    start_page_id,
-    end_page_id,
-    image_type,
-):
-    """供多进程调用"""
-    return load_images_from_pdf_core(
-        pdf_bytes,
-        dpi,
-        start_page_id,
-        end_page_id,
-        image_type,
-    )
-
-
-def load_images_from_pdf(
-    pdf_bytes: bytes,
-    dpi=200,
-    start_page_id=0,
-    end_page_id=None,
-    image_type=ImageType.PIL,
-    timeout=None,
-    threads=4,
-):
-    """
-    带超时控制的 PDF 转图片函数，支持多进程。
-
-    Args:
-        pdf_bytes (bytes): PDF bytes.
-        dpi (int): render dpi.
-        start_page_id (int): start page.
-        end_page_id (int | None): end page.
-        image_type (ImageType): PIL or BASE64.
-        timeout (int | None): timeout seconds.
-        threads (int): process number.
-    """
-
-    pdf_doc = fitz.open(stream=pdf_bytes, filetype="pdf")
-
-    if is_windows_environment() or CONFIG.PDF_TOOLS_WORKER_MAX_NUM==0:
-        return (
-            load_images_from_pdf_core(
-                pdf_bytes,
-                dpi,
-                start_page_id,
-                get_end_page_id(end_page_id, len(pdf_doc)),
-                image_type,
-            ),
-            pdf_doc,
-        )
-
-    if timeout is None:
-        timeout = get_load_images_timeout()
-
-    end_page_id = get_end_page_id(end_page_id, len(pdf_doc))
-
-    total_pages = end_page_id - start_page_id + 1
-
-    actual_threads = min(max(1, int(os.cpu_count() * CONFIG.PDF_TOOLS_WORKER_RATIO)), threads, total_pages)
-
-    pages_per_thread = max(1, total_pages // actual_threads)
-
-    page_ranges = []
-
-    for i in range(actual_threads):
-        range_start = start_page_id + i * pages_per_thread
-
-        if i == actual_threads - 1:
-            range_end = end_page_id
-        else:
-            range_end = start_page_id + (i + 1) * pages_per_thread - 1
-
-        page_ranges.append((range_start, range_end))
-
-
-    with ProcessPoolExecutor(max_workers=actual_threads) as executor:
-
-        futures = []
-
-        for range_start, range_end in page_ranges:
-            future = executor.submit(
-                _load_images_from_pdf_worker,
-                pdf_bytes,
-                dpi,
-                range_start,
-                range_end,
-                image_type,
-            )
-            futures.append((range_start, future))
-
-        try:
-            all_results = []
-
-            for range_start, future in futures:
-                images_list = future.result(timeout=timeout)
-                all_results.append((range_start, images_list))
-
-            all_results.sort(key=lambda x: x[0])
-
-            images_list = []
-
-            for _, imgs in all_results:
-                images_list.extend(imgs)
-
-            return images_list, pdf_doc
-
-        except FuturesTimeoutError:
-            pdf_doc.close()
-            executor.shutdown(wait=False, cancel_futures=True)
-            raise TimeoutError(
-                f"PDF to images conversion timeout after {timeout}s"
-            )
-
-def load_images_from_pdf_core(
-    pdf_bytes: bytes,
-    dpi=200,
-    start_page_id=0,
-    end_page_id=None,
-    image_type=ImageType.PIL,  # PIL or BASE64
-):
-    images_list = []
-
-    pdf_doc = fitz.open(stream=pdf_bytes, filetype="pdf")
-
-    pdf_page_num = len(pdf_doc)
-
-    end_page_id = get_end_page_id(end_page_id, pdf_page_num)
-
-    for index in range(start_page_id, end_page_id + 1):
-        # logger.debug(f"Converting page {index}/{pdf_page_num} to image")
-        page = pdf_doc[index]
-        image_dict = pdf_page_to_image(
-            page,
-            dpi=dpi,
-            image_type=image_type,
-        )
-        images_list.append(image_dict)
-
-    pdf_doc.close()
-
-    return images_list
-
-def cut_image(
-    bbox: tuple,
-    page_num: int,
-    page_pil_img,
-    return_path,
-    image_writer,
-    scale=2,
-):
-    filename = f"{page_num}_{int(bbox[0])}_{int(bbox[1])}_{int(bbox[2])}_{int(bbox[3])}"
-    rel_img_path = f"{return_path}_{filename}.jpeg" if return_path is not None else None
-    crop_img = get_crop_img(bbox, page_pil_img, scale=scale)
-    img_bytes = image_to_bytes(crop_img, image_format="JPEG")
-    image_writer.add_image(img_bytes, rel_img_path) 
-    return rel_img_path
-
-
-def get_crop_img(bbox: tuple, pil_img, scale=2):
-    scale_bbox = (
-        int(bbox[0] * scale),
-        int(bbox[1] * scale),
-        int(bbox[2] * scale),
-        int(bbox[3] * scale),
-    )
-    return pil_img.crop(scale_bbox)
-
-
-# def get_crop_np_img(bbox: tuple, input_img, scale=2):
-#     if isinstance(input_img, Image.Image):
-#         np_img = np.asarray(input_img)
-#     elif isinstance(input_img, np.ndarray):
-#         np_img = input_img
-#     else:
-#         raise ValueError("Input must be a pillow object or a numpy array.")
-
-#     scale_bbox = (
-#         int(bbox[0] * scale),
-#         int(bbox[1] * scale),
-#         int(bbox[2] * scale),
-#         int(bbox[3] * scale),
-#     )
-#     return np_img[scale_bbox[1] : scale_bbox[3], scale_bbox[0] : scale_bbox[2]]
-
-
-def images_bytes_to_pdf_bytes(image_bytes):
-    pdf_buffer = BytesIO()
-    image = Image.open(BytesIO(image_bytes))
-    image = ImageOps.exif_transpose(image) or image
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-    image.save(
-        pdf_buffer,
-        format="PDF",
-        # save_all=True
-    )
-    pdf_bytes = pdf_buffer.getvalue()
-    pdf_buffer.close()
-    return pdf_bytes
-
-
-def get_page_size(page):
-    rect = page.rect
-    w = rect.width
-    h = rect.height
-    return (w, h)
-    

NaviOCR/tools/pdf_image_tools_pdfium.py

@@ -1,268 +0,0 @@
-
-import os
-from io import BytesIO
-
-import numpy as np
-import pypdfium2 as pdfium
-import io
-from loguru import logger
-from PIL import Image, ImageOps
-
-import NaviOCR.config as CONFIG
-from NaviOCR.tools.check_sys_env import is_windows_environment
-from NaviOCR.tools.os_env_config import get_load_images_timeout
-from NaviOCR.tools.pdf_reader import image_to_b64str, image_to_bytes, page_to_image
-from NaviOCR.tools.enum_class import ImageType
-from NaviOCR.tools.hash_utils import str_sha256
-from NaviOCR.tools.pdf_page_id import get_end_page_id
-from concurrent.futures import ProcessPoolExecutor, TimeoutError as FuturesTimeoutError
-
-def convert_pdf_bytes_to_bytes(pdf_bytes, valid_single_page_ids=None):
-    pdf = pdfium.PdfDocument(pdf_bytes)
-    output_pdf = pdfium.PdfDocument.new()
-
-    try:
-        total_pages = len(pdf)
-        if not valid_single_page_ids:
-            valid_single_page_ids = list(range(total_pages))
-
-        valid_single_page_ids = [page_id for page_id in valid_single_page_ids if 0 <= page_id < total_pages]
-
-        for page_index in valid_single_page_ids:
-            try:
-                output_pdf.import_pages(pdf, pages=[page_index])
-            except Exception as page_error:
-                logger.warning(
-                    f"Failed to import page {page_index}: "
-                    f"{page_error}, skipping this page."
-                )
-
-        output_buffer = io.BytesIO()
-        output_pdf.save(output_buffer)
-        output_bytes = output_buffer.getvalue()
-
-    except Exception as e:
-        logger.warning(
-            f"Error in converting PDF bytes: {e}, "
-            f"Using original PDF bytes."
-        )
-        output_bytes = pdf_bytes
-
-    pdf.close()
-    output_pdf.close()
-
-    return output_bytes
-
-def pdf_page_to_image(page: pdfium.PdfPage, dpi=200, image_type=ImageType.PIL) -> dict:
-    """Convert pdfium.PdfDocument to image, Then convert the image to base64.
-
-    Args:
-        page (_type_): pdfium.PdfPage
-        dpi (int, optional): reset the dpi of dpi. Defaults to 200.
-        image_type (ImageType, optional): The type of image to return. Defaults to ImageType.PIL.
-
-    Returns:
-        dict:  {'img_base64': str, 'img_pil': pil_img, 'scale': float }
-    """
-    pil_img, scale = page_to_image(page, dpi=dpi)
-    image_dict = {
-        "scale": scale,
-    }
-    if image_type == ImageType.BASE64:
-        image_dict["img_base64"] = image_to_b64str(pil_img)
-    else:
-        image_dict["img_pil"] = pil_img
-
-    return image_dict
-
-
-def _load_images_from_pdf_worker(
-    pdf_bytes, dpi, start_page_id, end_page_id, image_type
-):
-    """用于进程池的包装函数"""
-    return load_images_from_pdf_core(
-        pdf_bytes, dpi, start_page_id, end_page_id, image_type
-    )
-
-
-def load_images_from_pdf(
-    pdf_bytes: bytes,
-    dpi=200,
-    start_page_id=0,
-    end_page_id=None,
-    image_type=ImageType.PIL,
-    timeout=None,
-    threads=4,
-):
-    """带超时控制的 PDF 转图片函数,支持多进程加速
-
-    Args:
-        pdf_bytes (bytes): PDF 文件的 bytes
-        dpi (int, optional): reset the dpi of dpi. Defaults to 200.
-        start_page_id (int, optional): 起始页码. Defaults to 0.
-        end_page_id (int | None, optional): 结束页码. Defaults to None.
-        image_type (ImageType, optional): 图片类型. Defaults to ImageType.PIL.
-        timeout (int | None, optional): 超时时间(秒)。如果为 None，则从环境变量 NaviOCR_PDF_LOAD_IMAGES_TIMEOUT 读取，若未设置则默认为 300 秒。
-        threads (int): 进程数,默认 4
-
-    Raises:
-        TimeoutError: 当转换超时时抛出
-    """
-    pdf_doc = pdfium.PdfDocument(pdf_bytes)
-    if is_windows_environment() or CONFIG.PDF_TOOLS_WORKER_MAX_NUM==0:
-        # Windows 环境下不使用多进程
-        return load_images_from_pdf_core(
-            pdf_bytes,
-            dpi,
-            start_page_id,
-            get_end_page_id(end_page_id, len(pdf_doc)),
-            image_type,
-        ), pdf_doc
-    else:
-        if timeout is None:
-            timeout = get_load_images_timeout()
-        end_page_id = get_end_page_id(end_page_id, len(pdf_doc))
-
-        # 计算总页数
-        total_pages = end_page_id - start_page_id + 1
-
-        # 实际使用的进程数不超过总页数
-        actual_threads = min(max(1, int(os.cpu_count() * CONFIG.PDF_TOOLS_WORKER_RATIO)), threads, total_pages)
-
-        # 根据实际进程数分组页面范围
-        pages_per_thread = max(1, total_pages // actual_threads)
-        page_ranges = []
-
-        for i in range(actual_threads):
-            range_start = start_page_id + i * pages_per_thread
-            if i == actual_threads - 1:
-                # 最后一个进程处理剩余所有页面
-                range_end = end_page_id
-            else:
-                range_end = start_page_id + (i + 1) * pages_per_thread - 1
-
-            page_ranges.append((range_start, range_end))
-
-
-        with ProcessPoolExecutor(max_workers=actual_threads) as executor:
-            # 提交所有任务
-            futures = []
-            for range_start, range_end in page_ranges:
-                future = executor.submit(
-                    _load_images_from_pdf_worker,
-                    pdf_bytes,
-                    dpi,
-                    range_start,
-                    range_end,
-                    image_type,
-                )
-                futures.append((range_start, future))
-
-            try:
-                # 收集结果并按页码排序
-                all_results = []
-                for range_start, future in futures:
-                    images_list = future.result(timeout=timeout)
-                    all_results.append((range_start, images_list))
-
-                # 按起始页码排序并合并结果
-                all_results.sort(key=lambda x: x[0])
-                images_list = []
-                for _, imgs in all_results:
-                    images_list.extend(imgs)
-
-                return images_list, pdf_doc
-            except FuturesTimeoutError:
-                pdf_doc.close()
-                executor.shutdown(wait=False, cancel_futures=True)
-                raise TimeoutError(f"PDF to images conversion timeout after {timeout}s")
-
-
-def load_images_from_pdf_core(
-    pdf_bytes: bytes,
-    dpi=200,
-    start_page_id=0,
-    end_page_id=None,
-    image_type=ImageType.PIL,  # PIL or BASE64
-):
-    images_list = []
-    pdf_doc = pdfium.PdfDocument(pdf_bytes)
-    pdf_page_num = len(pdf_doc)
-    end_page_id = get_end_page_id(end_page_id, pdf_page_num)
-
-    for index in range(start_page_id, end_page_id + 1):
-        # logger.debug(f"Converting page {index}/{pdf_page_num} to image")
-        page = pdf_doc[index]
-        image_dict = pdf_page_to_image(page, dpi=dpi, image_type=image_type)
-        images_list.append(image_dict)
-
-    pdf_doc.close()
-
-    return images_list
-
-
-def cut_image(
-    bbox: tuple,
-    page_num: int,
-    page_pil_img,
-    return_path,
-    image_writer,
-    scale=2,
-):
-    filename = f"{page_num}_{int(bbox[0])}_{int(bbox[1])}_{int(bbox[2])}_{int(bbox[3])}"
-    rel_img_path = f"{return_path}_{filename}.jpeg" if return_path is not None else None
-    crop_img = get_crop_img(bbox, page_pil_img, scale=scale)
-    img_bytes = image_to_bytes(crop_img, image_format="JPEG")
-    image_writer.add_image(img_bytes, rel_img_path) 
-    return rel_img_path
-
-
-def get_crop_img(bbox: tuple, pil_img, scale=2):
-    scale_bbox = (
-        int(bbox[0] * scale),
-        int(bbox[1] * scale),
-        int(bbox[2] * scale),
-        int(bbox[3] * scale),
-    )
-    return pil_img.crop(scale_bbox)
-
-
-# def get_crop_np_img(bbox: tuple, input_img, scale=2):
-#     if isinstance(input_img, Image.Image):
-#         np_img = np.asarray(input_img)
-#     elif isinstance(input_img, np.ndarray):
-#         np_img = input_img
-#     else:
-#         raise ValueError("Input must be a pillow object or a numpy array.")
-
-#     scale_bbox = (
-#         int(bbox[0] * scale),
-#         int(bbox[1] * scale),
-#         int(bbox[2] * scale),
-#         int(bbox[3] * scale),
-#     )
-
-#     return np_img[scale_bbox[1] : scale_bbox[3], scale_bbox[0] : scale_bbox[2]]
-
-
-def images_bytes_to_pdf_bytes(image_bytes):
-    pdf_buffer = BytesIO()
-    image = Image.open(BytesIO(image_bytes))
-    image = ImageOps.exif_transpose(image) or image
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-    # 第一张图保存为 PDF，其余追加
-    image.save(
-        pdf_buffer,
-        format="PDF",
-        # save_all=True
-    )
-    # 获取 PDF bytes 并重置指针（可选）
-    pdf_bytes = pdf_buffer.getvalue()
-    pdf_buffer.close()
-    return pdf_bytes
-
-def get_page_size(page):
-    w, h = page.get_size()
-    return (w, h)
-    

NaviOCR/tools/pdf_page_id.py

@@ -1,10 +0,0 @@
-
-from loguru import logger
-
-
-def get_end_page_id(end_page_id, pdf_page_num):
-    end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else pdf_page_num - 1
-    if end_page_id > pdf_page_num - 1:
-        logger.warning("end_page_id is out of range, use images length")
-        end_page_id = pdf_page_num - 1
-    return end_page_id

NaviOCR/tools/pdf_reader.py

@@ -1,111 +0,0 @@
-
-import base64
-from io import BytesIO
-
-from loguru import logger
-from PIL import Image
-import NaviOCR.config as CONFIG
-from pypdfium2 import PdfBitmap, PdfDocument, PdfPage 
-
-
-def page_to_image(
-    page: PdfPage,
-    dpi: int = 200,
-    max_width_or_height: int = 3500,  # changed from 4500 to 3500
-) -> (Image.Image, float):
-    scale = dpi / 72
-
-    long_side_length = max(*page.get_size())
-    if (long_side_length*scale) > max_width_or_height:
-        scale = max_width_or_height / long_side_length
-
-    bitmap: PdfBitmap = page.render(scale=scale)  # type: ignore
-
-    image = bitmap.to_pil()
-    try:
-        bitmap.close()
-    except Exception as e:
-        logger.error(f"Failed to close bitmap: {e}")
-    return image, scale
-
-
-def image_to_bytes(
-    image: Image.Image,
-    # image_format: str = "PNG",  # 也可以用 "JPEG"
-    image_format: str = "JPEG",
-) -> bytes:
-    with BytesIO() as image_buffer:
-        image.save(image_buffer, format=image_format)
-        return image_buffer.getvalue()
-
-
-def image_to_b64str(
-    image: Image.Image,
-    # image_format: str = "PNG",  # 也可以用 "JPEG"
-    image_format: str = "JPEG",
-) -> str:
-    image_bytes = image_to_bytes(image, image_format)
-    return base64.b64encode(image_bytes).decode("utf-8")
-
-
-def base64_to_pil_image(
-    base64_str: str,
-) -> Image.Image:
-    """Convert base64 string to PIL Image."""
-    image_bytes = base64.b64decode(base64_str)
-    with BytesIO(image_bytes) as image_buffer:
-        return Image.open(image_buffer).convert("RGB")
-
-
-def pdf_to_images(
-    pdf: str | bytes | PdfDocument,
-    dpi: int = 200,
-    max_width_or_height: int = 3500,
-    start_page_id: int = 0,
-    end_page_id: int | None = None,
-) -> list[Image.Image]:
-    doc = pdf if isinstance(pdf, PdfDocument) else PdfDocument(pdf)
-    page_num = len(doc)
-
-    end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else page_num - 1
-    if end_page_id > page_num - 1:
-        logger.warning("end_page_id is out of range, use images length")
-        end_page_id = page_num - 1
-
-    images = []
-    try:
-        for i in range(start_page_id, end_page_id + 1):
-            image, _ = page_to_image(doc[i], dpi, max_width_or_height)
-            images.append(image)
-    finally:
-        try:
-            doc.close()
-        except Exception:
-            pass
-    return images
-
-
-def pdf_to_images_bytes(
-    pdf: str | bytes | PdfDocument,
-    dpi: int = 200,
-    max_width_or_height: int = 3500,
-    start_page_id: int = 0,
-    end_page_id: int | None = None,
-    # image_format: str = "PNG",  # 也可以用 "JPEG"
-    image_format: str = "JPEG",
-) -> list[bytes]:
-    images = pdf_to_images(pdf, dpi, max_width_or_height, start_page_id, end_page_id)
-    return [image_to_bytes(image, image_format) for image in images]
-
-
-def pdf_to_images_b64strs(
-    pdf: str | bytes | PdfDocument,
-    dpi: int = 200,
-    max_width_or_height: int = 3500,
-    start_page_id: int = 0,
-    end_page_id: int | None = None,
-    # image_format: str = "PNG",  # 也可以用 "JPEG"
-    image_format: str = "JPEG",
-) -> list[str]:
-    images = pdf_to_images(pdf, dpi, max_width_or_height, start_page_id, end_page_id)
-    return [image_to_b64str(image, image_format) for image in images]

NaviOCR/tools/read_file.py

@@ -1,21 +0,0 @@
-from .guess_suffix_or_lang import guess_suffix_by_bytes
-from .pdf_image_tools import images_bytes_to_pdf_bytes
-from pathlib import Path
-
-pdf_suffixes = ["pdf"]
-image_suffixes = ["png", "jpeg", "jp2", "webp", "gif", "bmp", "jpg", "tiff"]
-
-
-
-def read_fn(path):
-    if not isinstance(path, Path):
-        path = Path(path)
-    with open(str(path), "rb") as input_file:
-        file_bytes = input_file.read()
-        file_suffix = guess_suffix_by_bytes(file_bytes, path)
-        if file_suffix in image_suffixes:
-            return images_bytes_to_pdf_bytes(file_bytes)
-        elif file_suffix in pdf_suffixes:
-            return file_bytes
-        else:
-            raise Exception(f"Unknown file suffix: {file_suffix}")

NaviOCR/version.py

@@ -1 +0,0 @@
-__version__='1.0.0'

NaviOCR/vlm_utils/NaviOCR_model.py

@@ -1,126 +0,0 @@
-import os
-from .NaviOCR_client import NaviOCRClient
-import NaviOCR.config as CONFIG
-
-
-class NaviOCRMODEL:
-    _instance = None
-    _models = {}
-    def __new__(cls, *args, **kwargs):
-        if cls._instance is None:
-            cls._instance = super().__new__(cls)
-        return cls._instance
-    
-    def get_model(
-        self,
-        backend: str,
-        model_path: str | None,
-        server_url: str | None,
-        **kwargs,
-    ) -> NaviOCRClient:
-        print(backend, model_path, server_url)
-        key = (backend, model_path, server_url)
-        print(key not in self._models)
-        if key not in self._models:
-            model = None
-            processor = None
-            vllm_llm = None
-            lmdeploy_engine = None
-            vllm_async_llm = None
-            batch_size = 0
-            max_concurrency = 100
-            http_timeout = 600
-            server_headers = None
-            max_retries = 3
-            retry_backoff_factor = 0.5
-            if os.getenv('OMP_NUM_THREADS') is None:
-                os.environ["OMP_NUM_THREADS"] = "1"
-            if backend == "transformers":
-                if not model_path:
-                    raise ValueError("model_path must be provided for the transformers backend.")
-                try:
-                    import torch
-                    from transformers import AutoProcessor
-                    try:
-                        from transformers import AutoModelForImageTextToText as AutoModelClass
-                    except ImportError:
-                        try:
-                            from transformers import AutoModelForVision2Seq as AutoModelClass
-                        except ImportError:
-                            from transformers import AutoModel as AutoModelClass
-                except ImportError as exc:
-                    raise ImportError("Please install torch and transformers to use the transformers backend.") from exc
-
-                processor = AutoProcessor.from_pretrained(
-                    model_path,
-                    trust_remote_code=True,
-                )
-                torch_dtype = kwargs.pop(
-                    "torch_dtype",
-                    torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-                )
-                model = AutoModelClass.from_pretrained(
-                    model_path,
-                    trust_remote_code=True,
-                    torch_dtype=torch_dtype,
-                    **kwargs,
-                )
-                if torch.cuda.is_available():
-                    model = model.cuda()
-                model = model.eval()
-            if backend == "vllm-engine":
-                try:
-                    import vllm
-                except ImportError:
-                    raise ImportError("Please install vllm to use the vllm-engine backend.")
-                if "gpu_memory_utilization" not in kwargs:
-                    kwargs["gpu_memory_utilization"] = CONFIG.GPU_MEMORY_UTILIZATION
-                if "model" not in kwargs:
-                    kwargs["model"] = model_path
-                if "logits_processors" not in kwargs:
-                    from . import LogitsProcessor
-                    kwargs["logits_processors"] = [LogitsProcessor]
-                if 'max_model_len' not in kwargs:
-                    kwargs["max_model_len"] = CONFIG.MAX_MODEL_LEN
-                vllm_llm = vllm.LLM(**kwargs)
-            
-            elif backend == "vllm-async-engine":
-                try:
-                    from vllm.engine.arg_utils import AsyncEngineArgs
-                    from vllm.v1.engine.async_llm import AsyncLLM
-                except ImportError:
-                    raise ImportError("Please install vllm to use the vllm-async-engine backend.")
-                
-                if "gpu_memory_utilization" not in kwargs:
-                    kwargs["gpu_memory_utilization"] = CONFIG.GPU_MEMORY_UTILIZATION
-                
-                if "model" not in kwargs:
-                    kwargs["model"] = model_path
-                
-                if "logits_processors" not in kwargs:
-                    from . import NaviOCRLogitsProcessor
-                    kwargs["logits_processors"] = [NaviOCRLogitsProcessor]
-            
-                if 'max_model_len' not in kwargs:
-                    kwargs["max_model_len"] = CONFIG.MAX_MODEL_LEN
-                vllm_async_llm = AsyncLLM.from_engine_args(AsyncEngineArgs(**kwargs))
-            
-            self._models[key] = NaviOCRClient(
-                backend=backend,
-                model=model,
-                processor=processor,
-                lmdeploy_engine=lmdeploy_engine,
-                vllm_llm=vllm_llm,
-                vllm_async_llm=vllm_async_llm,
-                server_url=server_url,
-                batch_size=batch_size,
-                max_concurrency=max_concurrency,
-                http_timeout=http_timeout,
-                server_headers=server_headers,
-                max_retries=max_retries,
-                retry_backoff_factor=retry_backoff_factor,
-            )
-        return self._models[key]
-
-
-NaviOCRMODEL_SERVICE = NaviOCRMODEL()

NaviOCR/vlm_utils/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.21"

NaviOCR/vlm_utils/vlm_client/vllm_async_engine_client.py

@@ -1,259 +0,0 @@
-import asyncio
-import uuid
-from io import BytesIO
-from typing import TYPE_CHECKING, Sequence
-
-if TYPE_CHECKING:
-    from vllm.outputs import RequestOutput
-
-from PIL import Image
-
-from .base_client import (
-    DEFAULT_SYSTEM_PROMPT,
-    DEFAULT_USER_PROMPT,
-    RequestError,
-    SamplingParams,
-    ServerError,
-    UnsupportedError,
-    VlmClient,
-)
-from .utils import aio_load_resource, gather_tasks, get_rgb_image
-
-
-class VllmAsyncEngineVlmClient(VlmClient):
-    def __init__(
-        self,
-        vllm_async_llm,  # vllm.v1.engine.async_llm.AsyncLLM instance
-        prompt: str = DEFAULT_USER_PROMPT,
-        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
-        sampling_params: SamplingParams | None = None,
-        text_before_image: bool = False,
-        allow_truncated_content: bool = False,
-        max_concurrency: int = 100,
-        debug: bool = False,
-    ):
-        super().__init__(
-            prompt=prompt,
-            system_prompt=system_prompt,
-            sampling_params=sampling_params,
-            text_before_image=text_before_image,
-            allow_truncated_content=allow_truncated_content,
-        )
-
-        try:
-            from vllm import SamplingParams
-            from vllm.sampling_params import RequestOutputKind
-            from vllm.v1.engine.async_llm import AsyncLLM
-        except ImportError:
-            raise ImportError("Please install vllm to use VllmEngineVlmClient.")
-
-        if not vllm_async_llm:
-            raise ValueError("vllm_async_llm is None.")
-        if not isinstance(vllm_async_llm, AsyncLLM):
-            raise ValueError(f"vllm_async_llm must be an instance of {AsyncLLM}")
-
-        self.vllm_async_llm = vllm_async_llm
-        if vllm_async_llm.tokenizer is None:
-            raise ValueError("vllm_async_llm.tokenizer is None.")
-
-        tokenizer = vllm_async_llm.tokenizer
-        if hasattr(tokenizer, "get_lora_tokenizer"):
-            tokenizer = tokenizer.get_lora_tokenizer()  # type: ignore
-
-        self.tokenizer = tokenizer
-        self.model_max_length = vllm_async_llm.model_config.max_model_len
-        self.VllmSamplingParams = SamplingParams
-        self.VllmRequestOutputKind = RequestOutputKind
-        self.max_concurrency = max_concurrency
-        self.debug = debug
-
-    def build_messages(self, prompt: str) -> list[dict]:
-        prompt = prompt or self.prompt
-        messages = []
-        if self.system_prompt:
-            messages.append({"role": "system", "content": self.system_prompt})
-        if "<image>" in prompt:
-            prompt_1, prompt_2 = prompt.split("<image>", 1)
-            user_messages = [
-                *([{"type": "text", "text": prompt_1}] if prompt_1.strip() else []),
-                {"type": "image"},
-                *([{"type": "text", "text": prompt_2}] if prompt_2.strip() else []),
-            ]
-        elif self.text_before_image:
-            user_messages = [
-                {"type": "text", "text": prompt},
-                {"type": "image"},
-            ]
-        else:  # image before text, which is the default behavior.
-            user_messages = [
-                {"type": "image"},
-                {"type": "text", "text": prompt},
-            ]
-        messages.append({"role": "user", "content": user_messages})
-        return messages
-
-    def build_vllm_sampling_params(self, sampling_params: SamplingParams | None):
-        sp = self.build_sampling_params(sampling_params)
-
-        vllm_sp_dict = {
-            "temperature": sp.temperature,
-            "top_p": sp.top_p,
-            "top_k": sp.top_k,
-            "presence_penalty": sp.presence_penalty,
-            "frequency_penalty": sp.frequency_penalty,
-            "repetition_penalty": sp.repetition_penalty,
-            # max_tokens should smaller than model max length
-            "max_tokens": sp.max_new_tokens if sp.max_new_tokens is not None else self.model_max_length,
-        }
-
-        if sp.no_repeat_ngram_size is not None:
-            vllm_sp_dict["extra_args"] = {
-                "no_repeat_ngram_size": sp.no_repeat_ngram_size,
-                "debug": self.debug,
-            }
-
-        return self.VllmSamplingParams(
-            **{k: v for k, v in vllm_sp_dict.items() if v is not None},
-            skip_special_tokens=False,
-            output_kind=self.VllmRequestOutputKind.FINAL_ONLY,
-        )
-
-    def get_output_content(self, output: "RequestOutput") -> str:
-        if not output.finished:
-            raise ServerError("The output generation was not finished.")
-
-        choices = output.outputs
-        if not (isinstance(choices, list) and choices):
-            raise ServerError("No choices found in the output.")
-
-        finish_reason = choices[0].finish_reason
-        if finish_reason is None:
-            raise ServerError("Finish reason is None in the output.")
-        if finish_reason == "length":
-            if not self.allow_truncated_content:
-                raise RequestError("The output was truncated due to length limit.")
-            else:
-                print("Warning: The output was truncated due to length limit.")
-        elif finish_reason != "stop":
-            raise RequestError(f"Unexpected finish reason: {finish_reason}")
-
-        return choices[0].text
-
-    def predict(
-        self,
-        image: Image.Image | bytes | str,
-        prompt: str = "",
-        sampling_params: SamplingParams | None = None,
-        priority: int | None = None,
-    ) -> str:
-        raise UnsupportedError(
-            "Synchronous predict() is not supported in vllm-async-engine VlmClient(backend). "
-            "Please use aio_predict() instead. If you intend to use synchronous client, "
-            "please use vllm-engine VlmClient(backend)."
-        )
-
-    def batch_predict(
-        self,
-        images: Sequence[Image.Image | bytes | str],
-        prompts: Sequence[str] | str = "",
-        sampling_params: Sequence[SamplingParams | None] | SamplingParams | None = None,
-        priority: Sequence[int | None] | int | None = None,
-    ) -> list[str]:
-        raise UnsupportedError(
-            "Synchronous batch_predict() is not supported in vllm-async-engine VlmClient(backend). "
-            "Please use aio_batch_predict() instead. If you intend to use synchronous client, "
-            "please use vllm-engine VlmClient(backend)."
-        )
-
-    async def aio_predict(
-        self,
-        image: Image.Image | bytes | str,
-        prompt: str = "",
-        sampling_params: SamplingParams | None = None,
-        priority: int | None = None,
-    ) -> str:
-        if isinstance(image, str):
-            image = await aio_load_resource(image)
-        if not isinstance(image, Image.Image):
-            image = Image.open(BytesIO(image))
-        image = get_rgb_image(image)
-
-        chat_prompt: str = self.tokenizer.apply_chat_template(
-            self.build_messages(prompt),  # type: ignore
-            tokenize=False,
-            add_generation_prompt=True,
-        )
-
-        vllm_sp = self.build_vllm_sampling_params(sampling_params)
-
-        generate_kwargs = {}
-        if priority is not None:
-            generate_kwargs["priority"] = priority
-
-        last_output = None
-        async for output in self.vllm_async_llm.generate(
-            prompt={"prompt": chat_prompt, "multi_modal_data": {"image": image}},
-            sampling_params=vllm_sp,
-            request_id=str(uuid.uuid4()),
-            **generate_kwargs,
-        ):
-            last_output = output
-
-        if last_output is None:  # this should not happen
-            raise ServerError("No output from the server.")
-        result = self.get_output_content(last_output)
-        # if 'table' in prompt:
-        #     print(result)
-        return result
-
-    async def aio_batch_predict(
-        self,
-        images: Sequence[Image.Image | bytes | str],
-        prompts: Sequence[str] | str = "",
-        sampling_params: Sequence[SamplingParams | None] | SamplingParams | None = None,
-        priority: Sequence[int | None] | int | None = None,
-        semaphore: asyncio.Semaphore | None = None,
-        use_tqdm=False,
-        tqdm_desc: str | None = None,
-    ) -> list[str]:
-        if isinstance(prompts, str):
-            prompts = [prompts] * len(images)
-        if not isinstance(sampling_params, Sequence):
-            sampling_params = [sampling_params] * len(images)
-        if not isinstance(priority, Sequence):
-            priority = [priority] * len(images)
-
-        assert len(prompts) == len(images), "Length of prompts and images must match."
-        assert len(sampling_params) == len(images), "Length of sampling_params and images must match."
-        assert len(priority) == len(images), "Length of priority and images must match."
-
-        if semaphore is None:
-            semaphore = asyncio.Semaphore(self.max_concurrency)
-
-        async def predict_with_semaphore(
-            image: Image.Image | bytes | str,
-            prompt: str,
-            sampling_params: SamplingParams | None,
-            priority: int | None,
-        ):
-            async with semaphore:
-                return await self.aio_predict(
-                    image=image,
-                    prompt=prompt,
-                    sampling_params=sampling_params,
-                    priority=priority,
-                )
-
-        return await gather_tasks(
-            tasks=[
-                predict_with_semaphore(*args)
-                for args in zip(
-                    images,
-                    prompts,
-                    sampling_params,
-                    priority,
-                )
-            ],
-            use_tqdm=use_tqdm,
-            tqdm_desc=tqdm_desc,
-        )

NaviOCR/vlm_utils/vlm_client/vllm_engine_client.py

@@ -1,255 +0,0 @@
-import asyncio
-from io import BytesIO
-from typing import TYPE_CHECKING, Sequence
-
-if TYPE_CHECKING:
-    from vllm.outputs import RequestOutput
-    from vllm.sampling_params import SamplingParams as VllmSamplingParams
-
-from PIL import Image
-
-from .base_client import (
-    DEFAULT_SYSTEM_PROMPT,
-    DEFAULT_USER_PROMPT,
-    RequestError,
-    SamplingParams, 
-    ServerError,
-    UnsupportedError,
-    VlmClient,
-)
-from .utils import get_rgb_image, load_resource
-
-
-class VllmEngineVlmClient(VlmClient):
-    def __init__(
-        self,
-        vllm_llm,  # vllm.LLM instance
-        prompt: str = DEFAULT_USER_PROMPT,
-        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
-        sampling_params: SamplingParams | None = None,
-        text_before_image: bool = False,
-        allow_truncated_content: bool = False,
-        batch_size: int = 0,
-        use_tqdm: bool = True,
-        debug: bool = False,
-    ):
-        super().__init__(
-            prompt=prompt,
-            system_prompt=system_prompt,
-            sampling_params=sampling_params,
-            text_before_image=text_before_image,
-            allow_truncated_content=allow_truncated_content,
-        )
-
-        try:
-            from vllm import LLM, SamplingParams
-        except ImportError:
-            raise ImportError("Please install vllm to use VllmEngineVlmClient.")
-
-        if not vllm_llm:
-            raise ValueError("vllm_llm is None.")
-        if not isinstance(vllm_llm, LLM):
-            raise ValueError("vllm_llm must be an instance of vllm.LLM.")
-
-        self.vllm_llm = vllm_llm
-        self.tokenizer = vllm_llm.get_tokenizer()
-        self.model_max_length = vllm_llm.llm_engine.model_config.max_model_len
-        self.VllmSamplingParams = SamplingParams
-        self.batch_size = batch_size
-        self.use_tqdm = use_tqdm
-        self.debug = debug
-
-    def build_messages(self, prompt: str) -> list[dict]:
-        prompt = prompt or self.prompt
-        messages = []
-        if self.system_prompt:
-            messages.append({"role": "system", "content": self.system_prompt})
-        if "<image>" in prompt:
-            prompt_1, prompt_2 = prompt.split("<image>", 1)
-            user_messages = [
-                *([{"type": "text", "text": prompt_1}] if prompt_1.strip() else []),
-                {"type": "image"},
-                *([{"type": "text", "text": prompt_2}] if prompt_2.strip() else []),
-            ]
-        elif self.text_before_image:
-            user_messages = [
-                {"type": "text", "text": prompt},
-                {"type": "image"},
-            ]
-        else:  # image before text, which is the default behavior.
-            user_messages = [
-                {"type": "image"},
-                {"type": "text", "text": prompt},
-            ]
-        messages.append({"role": "user", "content": user_messages})
-        return messages
-
-    def build_vllm_sampling_params(self, sampling_params: SamplingParams | None):
-        sp = self.build_sampling_params(sampling_params)
-
-        vllm_sp_dict = {
-            "temperature": sp.temperature,
-            "top_p": sp.top_p,
-            "top_k": sp.top_k,
-            "presence_penalty": sp.presence_penalty,
-            "frequency_penalty": sp.frequency_penalty,
-            "repetition_penalty": sp.repetition_penalty,
-            # max_tokens should smaller than model max length
-            "max_tokens": sp.max_new_tokens if sp.max_new_tokens is not None else self.model_max_length,
-        }
-
-        if sp.no_repeat_ngram_size is not None:
-            vllm_sp_dict["extra_args"] = {
-                "no_repeat_ngram_size": sp.no_repeat_ngram_size,
-                "debug": self.debug,
-            }
-
-        return self.VllmSamplingParams(
-            **{k: v for k, v in vllm_sp_dict.items() if v is not None},
-            skip_special_tokens=False,
-        )
-
-    def get_output_content(self, output: "RequestOutput") -> str:
-        if not output.finished:
-            raise ServerError("The output generation was not finished.")
-
-        choices = output.outputs
-        if not (isinstance(choices, list) and choices):
-            raise ServerError("No choices found in the output.")
-
-        finish_reason = choices[0].finish_reason
-        if finish_reason is None:
-            raise ServerError("Finish reason is None in the output.")
-        if finish_reason == "length":
-            if not self.allow_truncated_content:
-                raise RequestError("The output was truncated due to length limit.")
-            else:
-                print("Warning: The output was truncated due to length limit.")
-        elif finish_reason != "stop":
-            raise RequestError(f"Unexpected finish reason: {finish_reason}")
-
-        return choices[0].text
-
-    def predict(
-        self,
-        image: Image.Image | bytes | str,
-        prompt: str = "",
-        sampling_params: SamplingParams | None = None,
-        priority: int | None = None,
-    ) -> str:
-        return self.batch_predict(
-            [image],  # type: ignore
-            [prompt],
-            [sampling_params],
-        )[0]
-
-    def batch_predict(
-        self,
-        images: Sequence[Image.Image | bytes | str],
-        prompts: Sequence[str] | str = "",
-        sampling_params: Sequence[SamplingParams | None] | SamplingParams | None = None,
-        priority: Sequence[int | None] | int | None = None,
-    ) -> list[str]:
-        if not isinstance(prompts, str):
-            assert len(prompts) == len(images), "Length of prompts and images must match."
-        if isinstance(sampling_params, Sequence):
-            assert len(sampling_params) == len(images), "Length of sampling_params and images must match."
-        if isinstance(priority, Sequence):
-            assert len(priority) == len(images), "Length of priority and images must match."
-
-        image_objs: list[Image.Image] = []
-        for image in images:
-            if isinstance(image, str):
-                image = load_resource(image)
-            if not isinstance(image, Image.Image):
-                image = Image.open(BytesIO(image))
-            image = get_rgb_image(image)
-            image_objs.append(image)
-
-        if isinstance(prompts, str):
-            chat_prompts: list[str] = [
-                self.tokenizer.apply_chat_template(
-                    self.build_messages(prompts),  # type: ignore
-                    tokenize=False,
-                    add_generation_prompt=True,
-                )
-            ] * len(images)
-        else:  # isinstance(prompts, Sequence[str])
-            chat_prompts: list[str] = [
-                self.tokenizer.apply_chat_template(
-                    self.build_messages(prompt),  # type: ignore
-                    tokenize=False,
-                    add_generation_prompt=True,
-                )
-                for prompt in prompts
-            ]
-
-        if not isinstance(sampling_params, Sequence):
-            vllm_sp_list = [self.build_vllm_sampling_params(sampling_params)] * len(images)
-        else:
-            vllm_sp_list = [self.build_vllm_sampling_params(sp) for sp in sampling_params]
-
-        outputs = []
-        batch_size = self.batch_size if self.batch_size > 0 else len(images)
-        batch_size = max(1, batch_size)
-
-        for i in range(0, len(images), batch_size):
-            batch_image_objs = image_objs[i : i + batch_size]
-            batch_chat_prompts = chat_prompts[i : i + batch_size]
-            batch_sp_list = vllm_sp_list[i : i + batch_size]
-            batch_outputs = self._predict_one_batch(
-                batch_image_objs,
-                batch_chat_prompts,
-                batch_sp_list,
-            )
-            outputs.extend(batch_outputs)
-
-        return outputs
-
-    def _predict_one_batch(
-        self,
-        image_objs: list[Image.Image],
-        chat_prompts: list[str],
-        vllm_sampling_params: list["VllmSamplingParams"],
-    ):
-        vllm_prompts = [
-            {"prompt": chat_prompt, "multi_modal_data": {"image": image}}
-            for chat_prompt, image in zip(chat_prompts, image_objs)
-        ]
-
-        outputs = self.vllm_llm.generate(
-            prompts=vllm_prompts,  # type: ignore
-            sampling_params=vllm_sampling_params,
-            use_tqdm=self.use_tqdm,
-        )
-
-        return [self.get_output_content(output) for output in outputs]
-
-    async def aio_predict(
-        self,
-        image: Image.Image | bytes | str,
-        prompt: str = "",
-        sampling_params: SamplingParams | None = None,
-        priority: int | None = None,
-    ) -> str:
-        raise UnsupportedError(
-            "Asynchronous aio_predict() is not supported in vllm-engine VlmClient(backend). "
-            "Please use predict() instead. If you intend to use asynchronous client, "
-            "please use vllm-async-engine VlmClient(backend)."
-        )
-
-    async def aio_batch_predict(
-        self,
-        images: Sequence[Image.Image | bytes | str],
-        prompts: Sequence[str] | str = "",
-        sampling_params: Sequence[SamplingParams | None] | SamplingParams | None = None,
-        priority: Sequence[int | None] | int | None = None,
-        semaphore: asyncio.Semaphore | None = None,
-        use_tqdm=False,
-        tqdm_desc: str | None = None,
-    ) -> list[str]:
-        raise UnsupportedError(
-            "Asynchronous aio_batch_predict() is not supported in vllm-engine VlmClient(backend). "
-            "Please use batch_predict() instead. If you intend to use asynchronous client, "
-            "please use vllm-async-engine VlmClient(backend)."
-        )

NaviOCR/vlm_utils/vlm_client/vllm_v1_no_repeat_ngram.py

@@ -1,93 +0,0 @@
-# Logits Processor for vLLM V1 Engine.
-
-from typing import Any
-
-import torch
-from vllm.config import VllmConfig
-
-try:
-    from vllm.v1.sample.logits_processor.interface import (
-        BatchUpdate,
-        LogitsProcessor,
-        MoveDirectionality,
-    )
-except ImportError as e:
-    raise ImportError("Please install vllm>=0.10.1 to use this feature.") from e
-
-
-def _get_int_value(extra_args: dict[str, Any] | None, key: str) -> int | None:
-    if isinstance(extra_args, dict):
-        arg_value = extra_args.get(key)
-        if arg_value is not None:
-            try:
-                return int(arg_value)
-            except Exception:
-                pass
-    return None
-
-
-class VllmV1NoRepeatNGramLogitsProcessor(LogitsProcessor):
-    """
-    Prevents repeating the same n-gram of specified size in the output.
-    Inspired by Hugging Face's NoRepeatNGramLogitsProcessor.
-    Handled Extra Args:
-        no_repeat_ngram_size (int): Size of the n-gram to avoid repeating.
-    """
-
-    def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool):
-        # mapping: index -> (no_repeat_ngram_size, output_tok_ids, cached_ngrams)
-        self.req_info: dict[int, tuple[int, list[int], dict[tuple, list[int]]]] = {}
-
-    def is_argmax_invariant(self) -> bool:
-        return False
-
-    def update_state(self, batch_update: BatchUpdate | None) -> None:
-        if not batch_update:
-            return
-
-        for index in batch_update.removed:
-            self.req_info.pop(index, None)
-
-        for index, params, _, output_tok_ids in batch_update.added:
-            val = _get_int_value(params.extra_args, "no_repeat_ngram_size")
-            no_repeat_ngram_size = 0 if (val is None or val < 0) else val
-            if isinstance(params.extra_args, dict) and params.extra_args.get("debug"):
-                print(f"Request {index}: no_repeat_ngram_size = {no_repeat_ngram_size}")
-            self.req_info[index] = (no_repeat_ngram_size, output_tok_ids, {})
-
-        for a_index, b_index, direct in batch_update.moved:
-            a_info = self.req_info.pop(a_index, None)
-            b_info = self.req_info.pop(b_index, None)
-            if a_info is not None:
-                self.req_info[b_index] = a_info
-            if direct == MoveDirectionality.SWAP and b_info is not None:
-                self.req_info[a_index] = b_info
-
-    def apply(self, logits: torch.Tensor) -> torch.Tensor:
-        for index in range(len(logits)):
-            req_info = self.req_info.get(index)
-            if req_info is None:
-                continue
-            no_repeat_ngram_size, output_tok_ids, cached_ngrams = req_info
-            if no_repeat_ngram_size <= 0:
-                continue
-            # Skip if there are not enough tokens to form an n-gram
-            if len(output_tok_ids) < no_repeat_ngram_size:
-                continue
-
-            # Get the n-gram prefix (all but the last token)
-            prev_ngram = tuple(output_tok_ids[-no_repeat_ngram_size:-1])
-            last_token = output_tok_ids[-1]
-
-            # Store this n-gram occurrence
-            cached_ngrams.setdefault(prev_ngram, []).append(last_token)
-
-            # Get the next-token candidates to ban based on current prefix
-            current_prefix = tuple(output_tok_ids[-no_repeat_ngram_size + 1 :])
-            banned_tokens = cached_ngrams.get(current_prefix, [])
-
-            # Set the logits of banned tokens to negative infinity
-            for token in banned_tokens:
-                logits[index][token] = -float("inf")
-
-        return logits