fedihch's picture
Update README.md
04a8ff0
|
raw
history blame
1.17 kB
metadata
language:
  - spa
  - en
license: other
datasets:
  - Custom dataset

InvoiceReceiptClassifier is a fine-tuned LayoutLMv2 model that classifies a document to an invoice or receipt.

Quick start: using the raw model

from transformers import (
    AutoModelForSequenceClassification,
    LayoutLMv2FeatureExtractor,
    LayoutLMv2Tokenizer,
    LayoutLMv2Processor,
)
model = AutoModelForSequenceClassification.from_pretrained("fedihch/InvoiceReceiptClassifier")
feature_extractor = LayoutLMv2FeatureExtractor()
tokenizer = LayoutLMv2Tokenizer.from_pretrained("microsoft/layoutlmv2-base-uncased")
processor = LayoutLMv2Processor(feature_extractor, tokenizer)
from PIL import Image
input_img = Image.open("*****.jpg")
w, h = input_img.size
input_img = input_img.convert("RGB").resize((int(w * 600 / h), 600))
encoded_inputs = processor(input_img, return_tensors="pt")
for k, v in encoded_inputs.items():
    encoded_inputs[k] = v.to(model.device)
outputs = model(**encoded_inputs)
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
id2label = {0: "invoice", 1: "receipt"}
print(id2label[predicted_class_idx])