Jwalit/kyc-document-corner-detector

KYC Document Corner Detector

Lightweight document segmentation model trained on KYC documents (Aadhaar, PAN, passports, visas) from the Jwalit/moire-docs dataset.

Model Details

Property	Value
Architecture	MobileNetV3-Small encoder + upsampling decoder
Task	Binary segmentation (document vs background)
Training	CPU only, 8 epochs
Images	461 (391 train, 70 val)
Best Val IoU	74.79%
Model size	~10 MB
Labels	Self-supervised via OpenCV contour detection

How It Works

1. Input: Raw KYC document image (any size)
2. Segmentation: Model predicts binary mask of document region
3. Corner Detection: OpenCV contour extraction finds 4 corners from mask
4. Perspective Transform: Crops to document boundaries

Self-Supervised Label Generation

Labels are generated automatically using classical computer vision:

• Grayscale → Gaussian blur → Adaptive thresholding
• Morphological closing connects text regions
• Largest contour extraction → 4-corner approximation

No manual annotation required.

Files

File	Description
pytorch_model.bin	Trained model weights
config.json	Model configuration
inference_pipeline.py	Complete inference script (crop + rotate)
train_rotation_classifier.py	Script to train rotation classifier

Usage

import torch
from inference_pipeline import SegModel, predict_corners

model = SegModel()
model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
corners = predict_corners(model, "your_document.jpg")

Related Model

• Rotation Classifier: https://huggingface.co/Jwalit/kyc-document-rotation-classifier
  • 4-class classifier: 0°, 90°, 180°, 270°
  • Run train_rotation_classifier.py to train on your CPU

Pipeline

Raw Image → SegModel → Mask → Contours → 4 Corners → Crop
                                    ↓
                              RotModel → Classify Rotation → Correct Orientation