AI-405 LayoutLM Classification Task

Document-Classification/LayoutLM/.gitignore

@@ -0,0 +1,5 @@
+encoded_data.jsonl
+results
+report_documents
+other_documents
+bb_output

Document-Classification/LayoutLM/README.md

@@ -0,0 +1,90 @@
+# LayoutLMv3 Document Classification Workflow
+
+This experiment demonstrates the end-to-end workflow for document classification using LayoutLMv3. The steps include data preparation, model training, evaluation, and prediction.
+
+## 1. Install Dependencies
+
+Initialize and enter your virtual environment.
+
+```sh
+python -m venv .venv
+source .venv/bin/activate
+```
+
+Use `pip` to install dependencies.
+
+```sh
+pip install -r requirements
+```
+
+## 2. Populate Image Folders
+
+In this test, two folders were used:
+
+- `report_documents`: This contains the sample monthly report forms. If you don't have your own, you can look to the `Form-Generation` folder in this repo to generate a sample set. This experiement used 100 generated forms.
+- `other_documents`: Any assortment of other document types. We used the [FUNSD](https://guillaumejaume.github.io/FUNSD/) sample documents.
+
+    ```ts
+    @inproceedings{jaume2019,
+        title = {FUNSD: A Dataset for Form Understanding in Noisy Scanned Documents},
+        author = {Guillaume Jaume, Hazim Kemal Ekenel, Jean-Philippe Thiran},
+        booktitle = {Accepted to ICDAR-OST},
+        year = {2019}
+    }
+    ```
+
+## 3. Prepare Training Data
+
+```sh
+python prepare_training_data.py
+```
+
+Run the `prepare_training_data.py` script to generate a `jsonl` file with sample training data. This requires that the image folders are created and populated.
+
+By default, it will save the results in the `encoded_data.jsonl` file. This file can be deceptively large, and editors like VS Code may have difficulty opening it.
+
+Currently, this file is hardcoded to apply one type of label to each folder. Either the document is a monthly report or it falls under the `other` type.
+
+For consistency, this is using OCR results from the LayoutLMv3 model. Documentation suggests you should be able to provide your own OCR data, but attempts at this were not successful.
+
+## 4. Train the Model
+
+```sh
+python train_model.py
+```
+
+This script applies the training data from the previous step in order to fine-tune the LayoutLMv3 model. It uses a hard-coded 90/10 split to assign training and evaulation data.
+
+While only two labels are specified, real training data would ideally contain much more.
+
+Training was done on a local system. This seemed sufficient for this experiment, and the results were good. Adjust the batch size and epoch count if you find additional tweaking is necessary.
+
+It will output the results in a checkpoint folder within `/results`.
+
+## 5. Evaluate the Model
+
+```sh
+python evaluate_model.py <model_path>
+```
+
+Take note of where your previous checkpoint was saved. The `model_path` argument should refer to the folder, not a specific file. For example, if the checkpoint folder is `checkpoint-1`, then your command would be `python evaluate_model results/checkpoint-1`.
+
+The output should be a measurement of accuracy (1 is 100% accurate), along with other metrics.
+
+## 6. Use the Model
+
+```sh
+python use_model.py <model_path> <image_path>
+```
+
+Finally, you can check the predicted label for a specific document with `use_model.py`. Provide the model path and image path as arguments.
+
+The output will contain the predicted label for that document. As per the `labels_enum.py` file, the label is actually an integer value, with 0 representing a monthly report and 1 representing any other document.
+
+## 7. Demo with Streamlit
+
+There's a simple UI to demo your checkpoints.
+
+Run `streamlit run streamlit_app.py`, and it should start in your browser.
+
+In the UI, select your checkpoint and upload an image for the classification output.

Document-Classification/LayoutLM/evaluate_model.py

@@ -0,0 +1,45 @@
+import json
+from datasets import Dataset
+from transformers import (
+    LayoutLMv3ForSequenceClassification,
+    Trainer,
+)
+import sys
+
+if len(sys.argv) != 2:
+    print("Usage: python evaluate_model.py <model_path>")
+    sys.exit(1)
+
+model_path = sys.argv[1]
+
+# Load eval data
+with open("encoded_data.jsonl", "r") as f:
+    data = [json.loads(line) for line in f]
+
+# Shuffle and split (same as training)
+import random
+
+random.shuffle(data)
+
+eval_dataset = Dataset.from_list(data)
+
+# Load model
+model = LayoutLMv3ForSequenceClassification.from_pretrained(model_path)
+
+
+# Define compute_metrics for accuracy
+def compute_metrics(eval_pred):
+    logits, labels = eval_pred
+    preds = logits.argmax(-1)
+    acc = (preds == labels).astype(float).mean().item()
+    return {"accuracy": acc}
+
+
+# Evaluate
+trainer = Trainer(
+    model=model,
+    eval_dataset=eval_dataset,
+    compute_metrics=compute_metrics,
+)
+results = trainer.evaluate()
+print("Evaluation results:", results)

Document-Classification/LayoutLM/labels_enum.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class Label(Enum):
+    MONTHLY_REPORT = 0
+    OTHER = 1

Document-Classification/LayoutLM/prepare_data_pytesseract.py

@@ -0,0 +1,28 @@
+# An optional script that gathers ocr data from an image.
+# Couldn't get LayoutLM to accept this input, although it doesn't appear to be
+# a fault with the input itself.
+
+from pytesseract import image_to_data
+from PIL import Image
+import json
+
+image_path = "report_documents/test_form.jpg"
+image = Image.open(image_path)
+ocr_data = image_to_data(image, output_type="dict")
+
+tokens = ocr_data["text"]
+bboxes = [
+    [
+        ocr_data["left"][i],
+        ocr_data["top"][i],
+        ocr_data["left"][i] + ocr_data["width"][i],
+        ocr_data["top"][i] + ocr_data["height"][i],
+    ]
+    for i in range(len(ocr_data["text"]))
+]
+
+# Save output to JSON file in bb_output folder
+output_data = {"image": image_path, "tokens": tokens, "bboxes": bboxes}
+with open("bb_output/ocr_output.jsonl", "w") as f:
+    json.dump(output_data, f)
+print("OCR output saved to bb_output/ocr_output.jsonl")

Document-Classification/LayoutLM/prepare_training_data.py

@@ -0,0 +1,50 @@
+import json
+from PIL import Image
+from transformers import LayoutLMv3Processor
+from labels_enum import Label
+
+import os
+
+# Directory containing images
+report_dir = "report_documents"
+other_dir = "other_documents"
+output_jsonl = "encoded_data.jsonl"
+
+processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
+
+
+def pad_to_length(seq, length, pad_value):
+    return seq + [pad_value] * (length - len(seq))
+
+
+def encode_folder(folder_path, label, max_length=512):
+    for fname in os.listdir(folder_path):
+        fpath = os.path.join(folder_path, fname)
+        print(fpath)
+        # Only images
+        if not (
+            fname.lower().endswith(".jpg")
+            or fname.lower().endswith(".png")
+            or fname.lower().endswith(".jpeg")
+        ):
+            continue
+        # Image must have three dimensions (height, width, channels)
+        # Greyscale only has 2, so convert to guarantee the third
+        image = Image.open(fpath).convert("RGB")
+        encoding = processor(
+            image,
+            return_tensors="pt",
+            truncation=True,
+            padding="max_length",
+            max_length=max_length,
+        )
+        item = {}
+        item = {k: v[0].tolist() for k, v in encoding.items()}
+        item["image_path"] = fpath
+        item["label"] = label.value
+        out_f.write(json.dumps(item) + "\n")
+
+
+with open(output_jsonl, "w") as out_f:
+    encode_folder(other_dir, Label.OTHER)
+    encode_folder(report_dir, Label.MONTHLY_REPORT)

Document-Classification/LayoutLM/streamlit_app.py

@@ -0,0 +1,81 @@
+import streamlit as st
+from PIL import Image
+import io
+import os
+
+import torch
+from transformers import LayoutLMv3Processor, LayoutLMv3ForSequenceClassification
+from labels_enum import Label
+
+st.title("Document Classification Demo")
+
+
+# Sidebar for model selection
+st.sidebar.header("Model Settings")
+results_dir = "./results"
+checkpoints = [
+    d
+    for d in os.listdir(results_dir)
+    if d.startswith("checkpoint-") and os.path.isdir(os.path.join(results_dir, d))
+]
+checkpoints.sort()
+default_checkpoint = checkpoints[-1] if checkpoints else ""
+selected_checkpoint = st.sidebar.selectbox(
+    "Select model checkpoint",
+    checkpoints,
+    index=(
+        checkpoints.index(default_checkpoint)
+        if default_checkpoint in checkpoints
+        else 0
+    ),
+)
+model_path = (
+    os.path.join(results_dir, selected_checkpoint) if selected_checkpoint else None
+)
+
+
+# Load model and processor once per model_path
+@st.cache_resource
+def load_model(model_path):
+    processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
+    model = LayoutLMv3ForSequenceClassification.from_pretrained(model_path)
+    return processor, model
+
+
+processor, model = load_model(model_path)
+
+uploaded_file = st.file_uploader(
+    "Upload a document (PDF or image)", type=["png", "jpg", "jpeg"]
+)
+
+if uploaded_file:
+    file_bytes = uploaded_file.read()
+    file_name = uploaded_file.name
+    st.write(f"Uploaded: {file_name}")
+    # Only support image files for now
+    if file_name.lower().endswith((".png", ".jpg", ".jpeg")):
+        image = Image.open(io.BytesIO(file_bytes)).convert("RGB")
+        st.image(image, caption="Uploaded Image")
+        # Preprocess and run inference
+        inputs = processor(
+            image,
+            return_tensors="pt",
+            truncation=True,
+            padding="max_length",
+            max_length=512,
+        )
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+            pred = logits.argmax(-1).item()
+        # Map prediction to label name
+        label_name = None
+        for lbl in Label:
+            if lbl.value == pred:
+                label_name = lbl.name
+                break
+        st.success(f"Classification result: {label_name if label_name else pred}")
+    elif file_name.lower().endswith(".pdf"):
+        st.info("PDF preview not supported. Proceeding to classification.")
+else:
+    st.write("Please upload a document to classify.")

Document-Classification/LayoutLM/train_model.py

@@ -0,0 +1,41 @@
+# Load, shuffle, and split encoded data
+import json
+import random
+
+input_jsonl = "encoded_data.jsonl"
+with open(input_jsonl, "r") as f:
+    data = [json.loads(line) for line in f]
+
+# Don't want all the real documents together
+random.shuffle(data)
+
+# Use 90% of the data as training, the rest as evaulation
+split_idx = int(0.9 * len(data))
+train_data = data[:split_idx]
+eval_data = data[split_idx:]
+
+print(
+    f"Loaded {len(data)} samples. Split: {len(train_data)} train, {len(eval_data)} eval."
+)
+from transformers import LayoutLMv3ForSequenceClassification, Trainer, TrainingArguments
+from datasets import Dataset
+
+train_dataset = Dataset.from_list(train_data)
+eval_dataset = Dataset.from_list(eval_data)
+
+# This example only uses two labels, but in a larger scenario this should be based on training data.
+model = LayoutLMv3ForSequenceClassification.from_pretrained(
+    "microsoft/layoutlmv3-base", num_labels=2
+)
+
+training_args = TrainingArguments(
+    output_dir="./results", per_device_train_batch_size=4, num_train_epochs=3
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+)
+trainer.train()

Document-Classification/LayoutLM/use_model.py

@@ -0,0 +1,29 @@
+import sys
+from PIL import Image
+from transformers import LayoutLMv3Processor, LayoutLMv3ForSequenceClassification
+import torch
+
+if len(sys.argv) != 3:
+    print("Usage: python use_model.py <model_path> <image_path>")
+    sys.exit(1)
+
+model_path = sys.argv[1]
+image_path = sys.argv[2]
+
+# Load processor and model
+processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
+model = LayoutLMv3ForSequenceClassification.from_pretrained(model_path)
+
+# Load and preprocess image
+image = Image.open(image_path).convert("RGB")
+inputs = processor(
+    image, return_tensors="pt", truncation=True, padding="max_length", max_length=512
+)
+
+# Predict
+with torch.no_grad():
+    outputs = model(**inputs)
+    logits = outputs.logits
+    pred = logits.argmax(-1).item()
+
+print(f"Predicted label: {pred}")

Document-Classification/requirements

@@ -0,0 +1,6 @@
+transformers
+datasets
+torch
+torchvision
+accelerate
+streamlit