================================== 代码实现 ================================== PDF-Extract-Kit项目的核心代码实现在pdf_extract_kit目录下，该路径下包含下述几个模块： - configs: 特定模块的配置文件，如 ``pdf_extract_kit/configs/unimernet.yaml`` ，如果本身配置简单，建议放在 ``repo_root/configs`` 的 ``yaml`` 文件中的 ``model_config`` 里进行定义，方便用户修改。 - dataset: 自定义的 ``ImageDataset`` 类，用于加载和预处理图像数据。它支持多种输入类型，并且可以对图像进行统一的预处理操作（如调整大小、转换为张量等），以便于后续的模型推理加速。 - evaluation: 模型结果评测模块，支持多种任务类型评测，如 ``布局检测`` 、 ``公式检测`` 、 ``公式识别`` 等等，方便用户对不同任务、不同模型进行公平对比。 - registry: ``Registry`` 类是一个通用的注册表类，提供了注册、获取和列出注册项的功能。用户可以使用该类创建不同类型的注册表，例如任务注册表、模型注册表等。 - tasks: 最核心的任务模块，包含了许多不同类型的任务，如 ``布局检测`` 、 ``公式检测`` 、 ``公式识别`` 等等，用户添加新任务和新模型一般仅需要在这里进行代码添加。 .. note:: 基于上述的模块化设计，用户拓展新模块一般只需要在tasks里实现自己的新任务类及对应模型（更多情况下仅需要实现对应模型，任务已经定义好），然后在registry里注册即可。 下面我们以添加基于 ``YOLO``的 ``布局检测`` 模型为例，介绍如何添加新任务和新模型. 任务定义及注册 ============== 首先我们在 ``tasks`` 下添加一个 ``layout_detection`` 目录，然后在该目录下添加一个 ``task.py`` 文件用于定义布局检测任务类，具体如下： .. code-block:: python from pdf_extract_kit.registry.registry import TASK_REGISTRY from pdf_extract_kit.tasks.base_task import BaseTask @TASK_REGISTRY.register("layout_detection") class LayoutDetectionTask(BaseTask): def __init__(self, model): super().__init__(model) def predict_images(self, input_data, result_path): """ Predict layouts in images. Args: input_data (str): Path to a single image file or a directory containing image files. result_path (str): Path to save the prediction results. Returns: list: List of prediction results. """ images = self.load_images(input_data) # Perform detection return self.model.predict(images, result_path) def predict_pdfs(self, input_data, result_path): """ Predict layouts in PDF files. Args: input_data (str): Path to a single PDF file or a directory containing PDF files. result_path (str): Path to save the prediction results. Returns: list: List of prediction results. """ pdf_images = self.load_pdf_images(input_data) # Perform detection return self.model.predict(list(pdf_images.values()), result_path, list(pdf_images.keys())) 可以看到，任务定义包含下面几个要点： * 使用 ``@TASK_REGISTRY.register("layout_detection")`` 语法直接将布局任务类注册到 ``TASK_REGISTRY`` 下 ； * ``__init__`` 初始化函数传入 ``model`` , 具体参考 ``BaseTask`` 类 * 实现推理函数，这里考虑到布局检测通常会处理图像类及PDF文件，所以提供了两个函数 ``predict_images`` 和 ``predict_pdfs`` ，方便用户灵活选择。 模型定义及注册 ============== 接下来我们实现具体模型，在task下面新建models目录，并添加yolo.py用于YOLO模型定义，具体定义如下： .. code-block:: python import os import cv2 import torch from torch.utils.data import DataLoader, Dataset from ultralytics import YOLO from pdf_extract_kit.registry import MODEL_REGISTRY from pdf_extract_kit.utils.visualization import visualize_bbox from pdf_extract_kit.dataset.dataset import ImageDataset import torchvision.transforms as transforms @MODEL_REGISTRY.register('layout_detection_yolo') class LayoutDetectionYOLO: def __init__(self, config): """ Initialize the LayoutDetectionYOLO class. Args: config (dict): Configuration dictionary containing model parameters. """ # Mapping from class IDs to class names self.id_to_names = { 0: 'title', 1: 'plain text', 2: 'abandon', 3: 'figure', 4: 'figure_caption', 5: 'table', 6: 'table_caption', 7: 'table_footnote', 8: 'isolate_formula', 9: 'formula_caption' } # Load the YOLO model from the specified path self.model = YOLO(config['model_path']) # Set model parameters self.img_size = config.get('img_size', 1280) self.pdf_dpi = config.get('pdf_dpi', 200) self.conf_thres = config.get('conf_thres', 0.25) self.iou_thres = config.get('iou_thres', 0.45) self.visualize = config.get('visualize', False) self.device = config.get('device', 'cuda' if torch.cuda.is_available() else 'cpu') self.batch_size = config.get('batch_size', 1) def predict(self, images, result_path, image_ids=None): """ Predict layouts in images. Args: images (list): List of images to be predicted. result_path (str): Path to save the prediction results. image_ids (list, optional): List of image IDs corresponding to the images. Returns: list: List of prediction results. """ results = [] for idx, image in enumerate(images): result = self.model.predict(image, imgsz=self.img_size, conf=self.conf_thres, iou=self.iou_thres, verbose=False)[0] if self.visualize: if not os.path.exists(result_path): os.makedirs(result_path) boxes = result.__dict__['boxes'].xyxy classes = result.__dict__['boxes'].cls vis_result = visualize_bbox(image, boxes, classes, self.id_to_names) # Determine the base name of the image if image_ids: base_name = image_ids[idx] else: base_name = os.path.basename(image) result_name = f"{base_name}_MFD.png" # Save the visualized result cv2.imwrite(os.path.join(result_path, result_name), vis_result) results.append(result) return results 可以看到，模型定义包含下面几个要点： * 使用 ``@MODEL_REGISTRY.register('layout_detection_yolo')`` 语法直接将yolo布局模型注册到 ``MODEL_REGISTRY`` 下； * 初始化函数需要实现： + id_to_names的类别映射，用于可视化展示 + 模型参数配置 + 模型初始化 * 模型推理函数需要实现多种类型的模型推理：这里支持图像列表和PIL.Image类，可以方便用户直接基于图像路径或者图像流进行推理。 实现上述类定义后，将 ``LayoutDetectionYOLO`` 添加到 ``layout_detection`` 任务下 ``__init__.py`` 的 ``__all__`` 中即可。 .. code-block:: python from pdf_extract_kit.tasks.layout_detection.models.yolo import LayoutDetectionYOLO from pdf_extract_kit.registry.registry import MODEL_REGISTRY __all__ = [ "LayoutDetectionYOLO", ] .. note:: 对于同一个任务，我们支持多种模型，用户具体选择哪个可以根据评测结果进行选择，结合模型 ``精度`` 、 ``速度`` 和 ``场景适配程度`` 进行选择。 实现了任务和模型后，可以在 repo_root/scripts下添加脚本程序 ``layout_detection.py`` 示例脚本 ============== .. code-block:: python import os import sys import os.path as osp import argparse sys.path.append(osp.join(os.path.dirname(os.path.abspath(__file__)), '..')) from pdf_extract_kit.utils.config_loader import load_config, initialize_tasks_and_models import pdf_extract_kit.tasks # 确保所有任务模块被导入 TASK_NAME = 'layout_detection' def parse_args(): parser = argparse.ArgumentParser(description="Run a task with a given configuration file.") parser.add_argument('--config', type=str, required=True, help='Path to the configuration file.') return parser.parse_args() def main(config_path): config = load_config(config_path) task_instances = initialize_tasks_and_models(config) # get input and output path from config input_data = config.get('inputs', None) result_path = config.get('outputs', 'outputs'+'/'+TASK_NAME) # layout_detection_task model_layout_detection = task_instances[TASK_NAME] # for image detection detection_results = model_layout_detection.predict_images(input_data, result_path) # for pdf detection # detection_results = model_layout_detection.predict_pdfs(input_data, result_path) # print(detection_results) print(f'The predicted results can be found at {result_path}') if __name__ == "__main__": args = parse_args() main(args.config) 支持类型拓展 ============== 批处理拓展 ==============