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