llava_id_extraction

Multi-Head Vision Transformer for Indian Identity Document Classification

Document Classification Module of the LLaVA ID Extraction Project

This repository contains a fine-tuned Vision Transformer (ViT-Base) model that simultaneously predicts:

• 📄 Document Type
• 🪪 Document Side (Front / Back)
• 🗺️ Issuing State

using a shared Vision Transformer backbone with three independent classification heads.

This model serves as the first component of the larger LLaVA ID Extraction project, whose goal is to build an end-to-end AI pipeline for automated identity document understanding, information extraction, and verification.

---

Project Overview

Current Pipeline

Identity Document

        │

        ▼

Multi-Head Vision Transformer ✅

        │

        ├── Document Type

        ├── Front / Back

        └── State

        │

        ▼

OCR (Planned)

        │

        ▼

LLaVA Information Extraction (Planned)

        │

        ▼

Identity Verification (Planned)

        │

        ▼

Fraud Detection (Planned)

---

Model Details

Backbone

• Google ViT Base Patch16 224

Architecture

Shared Vision Transformer encoder with three task-specific classification heads:

1. Document Type Classification
2. Document Side Classification
3. State Classification

Framework

• PyTorch
• Hugging Face Transformers

---

Dataset

The model was trained on a synthetic dataset containing Indian identity document images with annotations for:

• Document Type
• Document Side
• Issuing State

The dataset includes multiple document categories and image variations suitable for supervised multi-task learning.

---

Training Configuration

---

Results

These results were obtained on the held-out validation split of the synthetic dataset.

---

Repository Contents

multihead_vit_best.pth
label_mappings.json
training.ipynb

---

Loading the Model

import torch

model.load_state_dict(
    torch.load(
        "multihead_vit_best.pth",
        map_location="cpu"
    )
)

model.eval()

---

Future Roadmap

This model represents the document classification module of the broader LLaVA ID Extraction project.

Future additions include:

• OCR integration
• LLaVA-based information extraction
• Identity verification
• Document forgery detection
• Fraud risk analysis
• Interactive Gradio demo

---

Limitations

The model was trained exclusively on synthetic identity document images.

Although it achieves excellent performance on the held-out synthetic validation set, additional fine-tuning and evaluation on real-world document scans and photographs are recommended before production deployment.

---

GitHub

GitHub repository:

https://github.com/IshankAggarwal09/llava_id_extraction