first-commit

classifier_app.py

@@ -0,0 +1,36 @@
+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]
+
+
+demo = gr.Interface(
+    fn=classify,
+    inputs=gr.Image(shape=(200, 200)),
+    outputs="text",
+    allow_flagging="manual",
+)
+demo.launch(share=True)