krishBSSPL/lipiocr-classifier

LipiOCR Document Classifier

A lightweight document-type classifier — the classification stage of LipiOCR, a fast document intelligence pipeline (classify → OCR → QR/barcode decode → normalize) built as a speed-focused complement to ZuraAI-VL, a larger VLM used for full structured extraction.

This model does one job well: given a photographed or scanned document, tell you what type of document it is, in milliseconds, on CPU or GPU. It does not read text — pair it with an OCR stage (LipiOCR uses glm-ocr) for full extraction.

Model details

• Architecture: EfficientNet-B0 (via timm), ImageNet-pretrained, fine-tuned end-to-end
• Input: 224×224 RGB image
• Output: 1 of 42 document type labels
• Parameters: ~5.3M

Categories (42)

Identity: aadhaar_card, cid, driving_license, nic, pan_card, passport, visa, voter_id Financial: bank_cheque, bank_statement, invoice, pos_payment_slip, purchase_order, quotation, receipt, telegraphic_transfer, payment_voucher, credit_card Travel: airway_bill, boarding_pass, courier_waybill Certificates: birth_certificate, certificate, marriage_certificate, medical_certificate, vaccination_certificate Business/other: business_card, employment_offer_letter, general_letter, income_tax_return, insurance_policy, marksheet, medical_prescription, menu_card, noc, product_catalog, reference_letter, rental_agreement, resume, salary_slip, utility_bill, vehicle_registration

Training data

• 7,227 training images across 42 categories, mostly synthetic (Faker-generated documents rendered with realistic layouts, colors, and — for the financial categories — real decodable QR/barcodes)
• A small set of real (non-synthetic) photographed/scanned documents, oversampled 15x during training to avoid being drowned out by the synthetic majority
• Training augmentation includes perspective warp, paper texture, directional lighting gradients, JPEG artifacts, and axis-aligned rotation — added specifically because a first version trained only on clean synthetic renders scored 0/11 on real photographed documents despite 98%+ synthetic validation accuracy. These fixes brought real-document accuracy to roughly 90%+ on most categories (see Limitations).

Performance

• 99.66% validation accuracy on a held-out synthetic split (1,167 examples)
• 11/12 real photographed documents correctly classified in a spot validation across categories with real examples available

Usage

import json
import torch
import timm
from PIL import Image
from torchvision import transforms
from safetensors.torch import load_file

config = json.load(open("config.json"))
model = timm.create_model(config["architecture"], pretrained=False, num_classes=config["num_classes"])
model.load_state_dict(load_file("model.safetensors"))
model.eval()

preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(config["normalize_mean"], config["normalize_std"]),
])

img = Image.open("document.jpg").convert("RGB")
# Recommended: run CLAHE contrast normalization + deskew on `img` before
# this step for best real-world accuracy - see preprocess.py in the
# LipiOCR repo. A plain resize works but skips that robustness step.
x = preprocess(img).unsqueeze(0)

with torch.no_grad():
    probs = torch.softmax(model(x), dim=1)[0]
top = torch.argmax(probs).item()
print(config["labels"][top], probs[top].item())

Limitations

• Marksheet vs. resume confusion: both render as structured line-item lists (subjects/grades vs. job history/skills), and this remains the weakest confusion pair in testing.
• Real-world academic transcripts (marksheets specifically) are the hardest category — real institutions' transcript layouts vary far more than any other document type, and the training set has limited real examples for this category specifically (~40% real-world accuracy vs. ~90%+ for most other categories).
• Trained primarily on English-language, Latin-script documents.
• No language identification or OCR capability — classification only.

Part of LipiOCR

This model is one stage of a larger pipeline. The full system (OCR extraction, QR/barcode decoding, digit-swap correction, and field normalization) lives in the LipiOCR project alongside this classifier.