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:
- Document Type Classification
- Document Side Classification
- 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
| Parameter | Value |
|---|---|
| GPU | NVIDIA Tesla T4 |
| Epochs | 4 |
| Optimizer | AdamW |
| Scheduler | Cosine Annealing |
| Mixed Precision | FP16 |
| Image Size | 224 × 224 |
| Backbone | ViT Base Patch16 224 |
Results
| Task | Accuracy |
|---|---|
| Document Type | 100% |
| Document Side | 100% |
| State | 100% |
| Combined Prediction | 100% |
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: