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)