navodPeiris/layoutlmv2-document-classifier

pip install PyMuPDF
pip install transformers
pip install torch
pip install torchvision
pip install pytesseract

import os
from PIL import Image
from transformers import LayoutLMv2Processor, LayoutLMv2ForSequenceClassification
import fitz
import io

processor = LayoutLMv2Processor.from_pretrained("microsoft/layoutlmv2-base-uncased")
model = LayoutLMv2ForSequenceClassification.from_pretrained("navodPeiris/layoutlmv2-document-classifier")

DATA_FOLDER = "data"
filename = "invoice.pdf"

file_location = os.path.join(DATA_FOLDER, filename)
doc = fitz.open(file_location)

page = doc.load_page(0)
pix = page.get_pixmap(dpi=200)

# Convert Pixmap to bytes
img_bytes = pix.tobytes("png")

# Load into PIL.Image
image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
doc.close()

encoding = processor(image, return_tensors="pt", truncation=True, padding="max_length", max_length=512)

outputs = model(**encoding)
logits = outputs.logits

predicted_class_id = logits.argmax(dim=1).item()
classified_output = model.config.id2label[predicted_class_id]

print(f"Predicted class: {classified_output}")