Início aplicação

app.py

@@ -0,0 +1,132 @@
+import io
+
+import pandas as pd
+import plotly.express as px
+import streamlit as st
+import torch
+import torch.nn.functional as F
+from easyocr import Reader
+from PIL import Image
+from transformers import (
+    LayoutLMv3FeatureExtractor,
+    LayoutLMv3TokenizerFast,
+    LayoutLMv3Processor,
+    LayoutLMv3ForSequenceClassification,
+)
+
+# DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
+DEVICE = "cpu"
+MICROSOFT_HODEL_NAME = "microsoft/layoutlmv3-base"
+MODEL_NAME = "arthur-lima/layoutlmv3-triagem-documentos"
+
+
+def create_bounding_box(bbox_data, width_scale: float, height_scale: float):
+    xs = []
+    ys = []
+    for x, y in bbox_data:
+        xs.append(x)
+        ys.append(y)
+    left = int(min(xs) * width_scale)
+    top = int(min(ys) * height_scale)
+    right = int(max(xs) * width_scale)
+    bottom = int(max(ys) * height_scale)
+    return [left, top, right, bottom]
+
+
+@st.experimental_singleton
+def create_ocr_reader():
+    # return Reader(["pt", "en"], gpu=True)
+    return Reader(["pt", "en"], gpu=False)
+
+
+@st.experimental_singleton
+def create_processor():
+    feature_extractor = LayoutLMv3FeatureExtractor(apply_ocr=False)
+    tokenizer = LayoutLMv3TokenizerFast.from_pretrained(MICROSOFT_HODEL_NAME)
+    return LayoutLMv3Processor(feature_extractor, tokenizer)
+
+
+@st.experimental_singleton
+def create_model():
+    model = LayoutLMv3ForSequenceClassification.from_pretrained(MODEL_NAME)
+    return model.eval().to(DEVICE)
+
+
+def predict(
+    image: Image.Image,
+    image_bytes: bytes,
+    reader: Reader,
+    processor: LayoutLMv3Processor,
+    model: LayoutLMv3ForSequenceClassification,
+):
+
+    ocr_result = reader.readtext(image_bytes)
+
+    width, height = image.size
+    width_scale = 1000 / width
+    height_scale = 1000 / height
+
+    words = []
+    boxes = []
+    for bbox, word, _ in ocr_result:
+        boxes.append(create_bounding_box(bbox, width_scale, height_scale))
+        words.append(word)
+
+    encoding = processor(
+        image,
+        words,
+        boxes=boxes,
+        max_length=512,
+        padding="max_length",
+        truncation=True,
+        return_tensors="pt",
+    )
+
+    with torch.inference_mode():
+        output = model(
+            input_ids=encoding["input_ids"].to(DEVICE),
+            attention_mask=encoding["attention_mask"].to(DEVICE),
+            bbox=encoding["bbox"].to(DEVICE),
+            pixel_values=encoding["pixel_values"].to(DEVICE),
+        )
+
+    logits = output.logits
+    predicted_class = logits.argmax()
+    probabilities = (
+        F.softmax(logits, dim=-1).flatten().tolist()
+    )  # Convertendo em probabilidades novamente
+    # return model.config.id2label[predicted_class.item()]
+    return predicted_class.detach().item(), probabilities
+
+
+reader = create_ocr_reader()
+processor = create_processor()
+model = create_model()
+uploaded_file = st.file_uploader("Upload Document Image", ["jpg", "png"])
+if uploaded_file is not None:
+    # Upload da imagem
+    image_bytes = uploaded_file.getvalue()
+    bytes_data = io.BytesIO(image_bytes)
+    image = Image.open(bytes_data)
+
+    # Mostrar a imagem
+    st.image(image, "Página do documento", width=300)
+
+    # Fazer a previsão
+    predicted_class, probabilities = predict(
+        image, image_bytes, reader, processor, model
+    )
+
+    # Imprimir o resultado na tela
+    predicted_label = model.config.id2label[predicted_class]
+    st.markdown(f"Tipo do documento previsto: **{predicted_label}**")
+
+    # Desenhar o gráfico de confianças
+    df_predictions = pd.DataFrame(
+        {
+            "Tipo Documento": list(model.config.id2label.values()),
+            "Confiança": probabilities,
+        }
+    )
+    fig = px.bar(df_predictions, x="Tipo Documento", y="Confiança")
+    st.plotly_chart(fig, use_container_width=True)

requirements.txt

@@ -0,0 +1,10 @@
+PyMuPDF==1.21.1
+numpy==1.24.2
+streamlit==1.15.2
+transformers==4.25.1
+pandas=2.0.0
+plotly-express=0.4.1
+python-dotenv==1.0.0
+Pillow=9.4.0
+torch=2.0.0
+easyocr=1.6.2