teddykoker / unsupervised-deep-homography

PyTorch implementation of Unsupervised Deep Homography: https://arxiv.org/abs/1709.03966
62 stars 22 forks source link

Unsupervised Deep Homography - Unoffical PyTorch Implementation

Unsupervised Deep Homography: A Fast and Robust Homography Estimation Model
Ty Nguyen, Steven W. Chen, Shreyas S. Shivakumar, Camillo J. Taylor, Vijay Kumar

models Figure from original paper. Proposed model is (c)

Implementation Details

This implementation leverages kornia, an open source differentiable computer vision library. Kornia is used for computing the direct linear transform (DLT) as well as the spatial transformation.

Uses PyTorch Lighting for easy GPU training and reproducibility.

model.py: Regression model implementation
dataset.py: Synthetic data generator
train.py: Train unsupervised model using photometric loss outlined in paper

Test

Download pre-trained weights

bash download_weights.sh

Create gifs:

python test.py path/to/test/images
Input Registered

Training

Note: tested on PyTorch version 1.4.0; previous versions have a bug that cause torch.inverse() and torch.solve() to generate runtime errors.

The model can be trained using synthetic data, created from the COCO dataset.

python train.py path/to/COCO/train/ path/to/COCO/valid/