---
license: apache-2.0
library_name: pytorch
pipeline_tag: image-classification

tags:
- vision-transformer
- pytorch
- computer-vision
- image-classification
- document-classification
- identity-documents
- multi-task-learning
- vit
---

# 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

| 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

```python
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**