Document Localization with Transfer Learning

Fine-tuning a pre-trained ResNet-18 to detect document boundaries in smartphone photos, enabling automated document rectification on mobile devices.

Project Context

This project was developed as a group final project for DS-301 (Deep Learning) at New York University, in collaboration with Sophie Juco during Fall 2022.

Motivation

Smartphone photos of documents frequently contain distortions — perspective skew, uneven lighting, and curvature — that reduce legibility and make automated processing difficult. Correcting these distortions requires first knowing where the document boundaries are. This project explores whether transfer learning on a pre-trained CNN can accurately localize document boundaries in distorted photos, as a first step toward automated rectification on mobile devices.

Dataset

SmartDoc 2015 is a document localization benchmark dataset containing 24,000 frames extracted from videos of 30 documents photographed on 4 different backgrounds, including challenging lighting conditions and partial occlusions.

We used the preprocessed version provided by Khurram Javed et al. (Recursive CNNs), stored in a Google Cloud Storage bucket and downloaded directly into the Colab training environment. The preprocessed dataset can be downloaded from this Google Drive folder. It contains 16,765 images in the training set and 6,158 images in the test set.

Ground truth boundaries are non-rectangular quadrilaterals (four corner coordinates), reflecting real-world perspective distortion. IoU was computed using a custom function built with Shapely to handle non-rectangular polygon overlap correctly.

Model Architecture

Base model: ResNet-18 pre-trained on ImageNet, sourced from qubvel/classification_models

Regression head: Global average pooling → Dropout (0.2) → Dense (8 units, ReLU), outputting 4 corner coordinates (x, y) pairs

Data augmentation: Random contrast and brightness adjustments applied during training

Training configuration:

Optimizer: Adam
Loss: Mean Squared Error (MSE)
Epochs: 10
Learning rate: 0.001 (head training), 0.00001 (fine-tuning)

Approach

Three training strategies were compared to evaluate the effect of progressively unlocking more of the pre-trained network:

1. Regression head training only: ResNet-18 weights frozen; only the regression head is trained.

2. Partial fine-tuning: The last 8 of 86 ResNet-18 layers unfrozen in addition to the regression head.

3. Full fine-tuning: All ResNet-18 layers unfrozen.

Results

Performance was evaluated using average Intersection-over-Union (IoU) across the test set.

Strategy	Avg. IoU	Training time	Prediction errors
Head training only	0.370	~639s/epoch	0 / 6,158
Partial fine-tuning	0.693	~678s/epoch	3 / 6,158
Full fine-tuning	0.822	~749s/epoch	0 / 6,158

Full fine-tuning achieved the best performance with only modest additional training time over partial fine-tuning. Partial fine-tuning produced a small number of invalid predictions (3 out of 6,158), likely due to instability introduced by partially unfreezing the network. The results confirm that deeper adaptation of the pre-trained weights is beneficial for this localization task.

Regression head training only

Example prediction (test image 0); IoU = 0.207:

Partial fine-tuning (8 layers)

Example prediction (test image 0); IoU = 0.723:

Prediction errors

Full fine-tuning

Example prediction (test image 0); IoU = 0.814:

Limitations & Next Steps

Partial fine-tuning instability: the source of invalid predictions is not fully understood; further investigation could make this a more practical middle-ground option
Model size vs. mobile deployment: ResNet-18 was chosen for its relatively low computational cost, but alternatives such as ResNet-50 or custom lightweight architectures could be explored for a better IoU/memory trade-off
Corner refinement: a separate CNN trained specifically to refine the predicted corner positions could improve boundary precision, particularly on high-distortion images

Repository Structure

File / Folder	Description
`training.ipynb`	Loads the dataset and trains the CNN under all three strategies
`prediction_analysis.ipynb`	Evaluates model performance using IoU
`utils.py`	Helper functions for analysis and visualization
`csv-cleaning.ipynb`	Removes unnecessary characters from ground truth files
`GDrive_to_GCS.ipynb`	Transfers the dataset from Google Drive to Google Cloud Storage
`Prediction-plots/`	Sample prediction results from each trained model
`Training-plots/`	Training and validation accuracy/loss curves for each strategy

Requirements

This project was developed on Google Colab in 2022 (TensorFlow 2.x assumed). Dependency versions are not pinned as the original environment is no longer available; adjustments may be needed for newer environments.

numpy
opencv-python
shapely
matplotlib
tensorflow

Note: Data loading code assumes access to a Google Cloud Storage bucket. Running locally requires adapting the data pipeline to load from disk.

Running the Project

This project was developed and run entirely on Google Colab. To reproduce:

Upload the notebooks to Google Colab
Download the SmartDoc 2015 dataset (preprocessed version by Javed et al.) and store it in a GCS bucket using GDrive_to_GSC.ipynb
Update the bucket path in the data loading code
Run csv_cleaning.ipynb to clean the ground-truth files for easier parsing
Load and fine-tune the model with training.ipynb
Analyze results with prediction_analysis.ipynb

(A local runner script is not currently provided.)