Real‑Time Document Corner Detection on Mobile: Heatmap‑Based Keypoint Algorithms Explained

Overview

Quark's end‑side intelligent group is working on real‑time document detection: given an RGB image, predict the four corner keypoints. The pipeline belongs to keypoint detection, so recent papers were reviewed and experiments performed.

Pipeline Decomposition

Image processing: optical augmentation, transformation, resize, crop to increase diversity.

Encoding: convert coordinates to labels for supervision.

Network model: backbone, FPN, detection head, etc.

Decoding: transform model output into Cartesian coordinates.

Related Works

Two main technical schemes for keypoint detection:

Face‑detection‑style: output tensor passes through a fully‑connected layer to obtain normalized 1‑D coordinates.

Human‑pose‑estimation‑style: output heatmaps, locate the maximum response and map back to image coordinates.

Heatmap‑based methods are dominant because they achieve better performance. The following recent papers are referenced: DSNT (2018), Distribution‑Aware Coordinate Representation (2019), Unbiased Data Processing (2019), AID (2020), etc.

Proposed Method

The authors combine the advantages of end‑to‑end optimization and spatial generalization by proposing a differentiable decoding method.

1. Idea

Two ways to obtain coordinates from heatmaps:

Argmax on heatmap → coordinate (good spatial generalization but non‑differentiable during training).

Heatmap → fully‑connected layer → coordinate (differentiable but loses spatial information).

The proposed method keeps the heatmap as a probability density, normalizes it, and computes the expected value using X and Y index matrices, which is fully differentiable and yields lower theoretical error.

2. Specific Steps

Model outputs K heatmaps of size H×W.

Normalize each heatmap to sum to 1 (treated as a discrete probability distribution).

Generate X and Y index matrices ranging from –1 to 1.

Compute the expected X and Y by matrix multiplication of the normalized heatmap with the index matrices.

3. Loss Function

The DSNT module uses a combination of Euclidean loss for coordinate regression and a JS divergence regularizer to force the heatmaps toward Gaussian distributions.

4. Advantages

End‑to‑end training aligns loss with test metrics.

Theoretical error lower.

Introducing X and Y matrices provides prior knowledge, easing learning.

Works well on low‑resolution inputs.

5. Limitations

Performance degrades when keypoints lie near image borders.

Unbiased Coordinate Transformations

Standard decoding suffers from quantization errors and misalignment after image flipping. The authors propose using unit length instead of image size to align heatmaps, and a combined classification‑regression format inspired by anchor‑based detection.

Additional Related Works

DSNT, AID (information dropping augmentation), RSN (multi‑person pose), Lite‑HRNet (lightweight high‑resolution network) are briefly described.

Model Architecture for Mobile

The final model adopts MobileNet‑v3‑small as backbone, FPN with nearest up‑sample + conv + BN + ReLU, and three branches (keypoints, mask, center). Only the keypoints branch is used at inference. Various optimization strategies were applied: mask and center auxiliary tasks, deep supervision, padding to avoid border issues, data augmentation (random crop, erase, flip), and loss function experiments (Euclidean, L1, L2, SmoothL1 – the latter performed best).

Evaluation Metrics

MSE for validation.

OKS‑mAP for keypoint similarity.

Inference time on Redmi 8 using MNN.

Experimental Results

A baseline (MobileNet‑v3 + FPN + SSH + keypoints + DSNT) was built without the optimization tricks. Subsequent experiments replaced loss functions and added tricks, reporting improvements and some ineffective attempts (edge branch, extra points).

Conclusion

For on‑device document keypoint detection, a heatmap + DSNT pipeline currently yields the best results, though OKS‑mAP still has room for improvement. Compared with FC‑based regression, heatmap methods cannot predict points outside the document region, which may cause missing content; future work should address this limitation.