import gradio as gr def classify(input_img): 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) 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"} return id2label[predicted_class_idx] examples =[['https://upserve.com/media/sites/2/Bill-from-Mezcalero-in-Washington-D.C.-photo-by-Alfredo-Solis-1-e1507226752437.jpg'], ['https://templates.invoicehome.com/invoice-template-us-neat-750px.png'],] demo = gr.Interface( fn=classify, inputs=gr.Image(shape=(200, 200)), outputs="text", allow_flagging="manual", description="Upload an invoice or receipt image and the model will classify it!", examples=examples, ) demo.launch(share=True)