DocumentClassifier ONNX

DS4SD DocumentClassifier 的高性能 ONNX 优化实现，用于高效的文档类型分类。

🎯 概述

DocumentClassifier 是一个专为自动文档类型分类设计的深度学习模型。此 ONNX 版本通过 JPQD（联合剪枝、量化和蒸馏）优化，为生产环境提供了增强的性能和优化推理能力。

主要特性

• 高准确率：在多个类别上实现可靠的文档类型分类
• 快速推理：CPU 上每份文档约 28ms（35+ FPS）
• 生产就绪：ONNX 格式支持跨平台部署
• 内存高效：通过 JPQD 压缩优化模型大小
• 易于集成：简洁的 Python API 和丰富的示例

🚀 快速开始

安装

pip install onnxruntime opencv-python pillow numpy

基本用法

from example import DocumentClassifierONNX
import cv2

# 初始化模型
classifier = DocumentClassifierONNX("DocumentClassifier.onnx")

# 从图像文件分类文档
result = classifier.classify("document.jpg")
print(f"Document type: {result['predicted_category']}")
print(f"Confidence: {result['confidence']:.3f}")

# 获取 top 预测结果
for pred in result['top_predictions']:
    print(f"{pred['category']}: {pred['confidence']:.3f}")

命令行界面

# 分类文档图像
python example.py --image document.jpg

# 运行性能基准测试
python example.py --benchmark --iterations 100

# 使用虚拟数据进行演示
python example.py

📊 模型规格

🏷️ 支持的文档类别

该模型可将文档分类为以下类别：

• Article（文章） - 新闻文章、博客帖子、网页内容
• Form（表单） - 申请表、调查问卷
• Letter（信函） - 商务信函、通信
• Memo（备忘录） - 内部备忘录、通知
• News（新闻） - 报纸文章、新闻稿
• Presentation（演示文稿） - 幻灯片、演示材料
• Resume（简历） - 简历、职业档案
• Scientific（学术文献） - 研究论文、学术文档
• Specification（技术规格） - 技术规范、手册
• Table（表格） - 数据表、电子表格内容
• Other（其他） - 其他文档类型

⚡ 性能基准

推理速度（CPU）

• 平均值：28.1ms ± 0.5ms
• 吞吐量：~35.6 FPS
• 硬件：现代 CPU（单线程）
• 批大小：1

内存使用

• 模型加载：~50MB RAM
• 推理时：~100MB RAM
• 峰值使用：~150MB RAM

🔧 高级用法

批量处理

import numpy as np
from example import DocumentClassifierONNX

classifier = DocumentClassifierONNX()

# 处理多个图像
image_paths = ["doc1.jpg", "doc2.pdf", "doc3.png"]
results = []

for path in image_paths:
    result = classifier.classify(path)
    results.append({
        'file': path,
        'category': result['predicted_category'],
        'confidence': result['confidence']
    })

# 显示结果
for r in results:
    print(f"{r['file']}: {r['category']} ({r['confidence']:.3f})")

自定义预处理

import cv2
import numpy as np

# 手动加载和预处理图像
image = cv2.imread("document.jpg")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 调整为模型输入大小
resized = cv2.resize(image, (224, 224))
normalized = resized.astype(np.float32) / 255.0

# 转换为 CHW 格式并添加批维度
chw = np.transpose(normalized, (2, 0, 1))
batched = np.expand_dims(chw, axis=0)

# 运行推理
classifier = DocumentClassifierONNX()
logits = classifier.predict(batched)
result = classifier.decode_output(logits)

🛠️ 集成示例

Flask Web 服务

from flask import Flask, request, jsonify
from example import DocumentClassifierONNX

app = Flask(__name__)
classifier = DocumentClassifierONNX()

@app.route('/classify', methods=['POST'])
def classify_document():
    file = request.files['document']
    
    # 保存并处理文件
    file.save('temp_document.jpg')
    result = classifier.classify('temp_document.jpg')
    
    return jsonify({
        'category': result['predicted_category'],
        'confidence': float(result['confidence']),
        'top_predictions': result['top_predictions']
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

批量处理脚本

import os
import glob
from example import DocumentClassifierONNX

def classify_directory(input_dir, output_file):
    classifier = DocumentClassifierONNX()
    
    # 查找所有图像文件
    extensions = ['*.jpg', '*.jpeg', '*.png', '*.pdf']
    files = []
    for ext in extensions:
        files.extend(glob.glob(os.path.join(input_dir, ext)))
    
    results = []
    for file_path in files:
        try:
            result = classifier.classify(file_path)
            results.append({
                'file': os.path.basename(file_path),
                'category': result['predicted_category'],
                'confidence': result['confidence']
            })
            print(f"✓ {file_path}: {result['predicted_category']}")
        except Exception as e:
            print(f"✗ {file_path}: Error - {e}")
    
    # 保存结果
    import json
    with open(output_file, 'w') as f:
        json.dump(results, f, indent=2)

# 用法
classify_directory("./documents", "classification_results.json")

📋 系统要求

系统要求

• Python：3.8 或更高版本
• 内存：至少 2GB 可用空间
• CPU：推荐 x86_64 架构
• 操作系统：Windows、Linux、macOS

依赖项

onnxruntime>=1.15.0
opencv-python>=4.5.0
numpy>=1.21.0
Pillow>=8.0.0

🔍 故障排除

常见问题

模型加载错误

# 确保模型文件存在
import os
if not os.path.exists("DocumentClassifier.onnx"):
    print("Model file not found!")

内存问题

# 对于低内存系统，逐个处理图像
# 并在使用后清理变量
import gc
result = classifier.classify(image)
del image  # 释放内存
gc.collect()

图像格式问题

# 将任何图像格式转换为 RGB
from PIL import Image
img = Image.open("document.pdf").convert("RGB")
result = classifier.classify(np.array(img))

📖 技术细节

架构

• 基础模型：深度卷积神经网络
• 输入处理：标准 ImageNet 预处理
• 特征提取：带有全局池化的 CNN 骨干网络
• 分类头：带有 softmax 激活的稠密层
• 优化：JPQD 量化以优化大小和速度

预处理流程

1. 图像加载：PIL/OpenCV 图像加载
2. 调整大小：双线性插值至 224×224
3. 归一化：[0, 255] → [0, 1] 范围
4. 格式转换：HWC → CHW（通道优先）
5. 添加批次：单张图像 → 批维度

输出处理

1. 特征提取：CNN 骨干网络输出 [1, 1280, 7, 7]
2. 全局池化：空间平均至 [1, 1280]
3. 分类：将特征映射到类别概率
4. Top-K 选择：返回最可能的类别

📚 引用

如果您在研究中使用此模型，请引用：

@article{docling2024,
  title={Docling Technical Report},
  author={DS4SD Team},
  journal={arXiv preprint arXiv:2408.09869},
  year={2024}
}

📄 许可证

本项目采用 MIT 许可证 - 详情请参见 LICENSE 文件。

🤝 贡献

欢迎贡献！请随时提交 Pull Request。对于重大更改，请先开 issue 讨论您想要更改的内容。

🆘 支持

• 问题反馈：GitHub Issues
• 文档：本 README 和代码内注释
• 示例：请参阅 example.py 获取完整的使用示例

📈 更新日志

v1.0.0

• 初始 ONNX 模型发布
• 应用 JPQD 优化
• 完整的 Python API
• CLI 命令行界面
• 完整的文档
• 性能基准测试

由 DS4SD 社区用 ❤️ 制作