项目介绍:dit-base-finetuned-rvlcdip
项目背景
“dit-base-finetuned-rvlcdip”是微软在Hugging Face平台发布的一个关于文档图像分类的项目。该项目使用了Document Image Transformer (DiT)模型,这是一种基于变换器(Transformer)的模型。基础模型首先在IIT-CDIP数据集上进行了自监督预训练,然后在RVL-CDIP数据集上进行了微调。IIT-CDIP是一个包含4200万张文档图像的数据集,而RVL-CDIP则是一个由40万张灰度图像组成的数据集,覆盖了16个类别,每个类别有25000张图片。
模型描述
Document Image Transformer(DiT)是一种类似BERT的变换器编码器模型。这个模型通过自监督学习的方式进行预训练,预训练的目标是通过dVAE(离散变分自编码器)的编码器预测被遮盖的视觉标记。图像会被以固定大小的片段(每个片段是16x16分辨率)形式输入模型,并通过线性嵌入后,再通过绝对位置嵌入,以序列的形式输入变换器编码器的各层。
通过这样的预训练步骤,模型能够学习图像的内部表示。这些表示可以用于特征提取,对后续任务有很大的帮助。例如,如果有一个带标签的文档图像数据集,可以在预训练编码器的基础上添加线性层来训练标准分类器。
预期用途及限制
这个模型可以直接用于将文档图像编码成向量空间,但主要用途是经过微调后用于诸如文档图像分类、表格检测、文档布局分析等任务。对于有兴趣的任务,可以访问模型库查找相应的微调版本。
如何使用
以下是在PyTorch中使用此模型的简单示例:
from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
from PIL import Image
# 打开并转换图像文件
image = Image.open('path_to_your_document_image').convert('RGB')
# 加载处理器和模型
processor = AutoImageProcessor.from_pretrained("microsoft/dit-base-finetuned-rvlcdip")
model = AutoModelForImageClassification.from_pretrained("microsoft/dit-base-finetuned-rvlcdip")
# 处理图像并得到输出
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
# 预测结果属于RVL-CDIP的16个类别之一
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])
使用这个预训练并微调的模型,可以大大提升处理文档图像的效率和准确度。通过合理的设置和训练,这个模型能够在实际应用中解决复杂的文档分类和分析任务。