cavalli1234/nefsac

Code for the ECCV 2022 paper NeFSAC: Neurally Filtered Minimal Samples

NeFSAC: Neurally Filtered Minimal Samples (ECCV 2022)

This repository provides tools to use NeFSAC, as described in this paper at ECCV 2022.

What you will find in this repository:

• Pretrained models for NeFSAC for Essential matrix estimation and Fundamental matrix estimation on both KITTI and PhotoTourism.
• Sample code to train NeFSAC on KITTI.
• A simple pytorch-only RANSAC implementation to showcase how to integrate NeFSAC into any RANSAC.

What is NeFSAC

NeFSAC is a lightweight model which is trained to recognize good minimal samples in RANSAC.

It learns common motion and depth patterns of correct correspondences in a dataset, and uses this prior to estimate the likelihood that a minimal set of correspondences could lead to a good model estimate. This knowledge can then be used in RANSAC to process the promising minimal samples first, and terminate much earlier. Moreover, it learns to recognize and avoid the degenerate patterns which lead to unstable model estimates, thus improving the accuracy of estimation on average.

In our experiments, we consistently observed improvements in accuracy in pose estimation (both E and F matrix), speed-ups between 2x and 10x on a good USAC implementation, and a great generalization capability. More details are in the paper.

How to run the NeFSAC demo

You can run NeFSAC on our simple pytorch-RANSAC on some sample images. You only need to install dependencies, and then run:

python demo.py

If your torch installation does not support CUDA, you can still run:

python demo.py --device cpu

You can also specify other image pairs by specifying parameters --image1 and --image2.

Note that this is a very simple baseline to showcase how NeFSAC is used in a RANSAC context. Best results are obtained with state-of-the-art RANSAC implementations. We are working on making available a NeFSAC sampler in MAGSAC++, stay tuned!

How to train NeFSAC

We observed impressive generalization capabilities with NeFSAC, so re-training on a new dataset might not be necessary. Nonetheless, we provide our training procedure in a Pytorch Lightning module in source/trainer.py, and show how data is preprocessed in source/data.py.

We provide code for reproducing training on KITTI. You should first download the KITTI odometry dataset in a local directory <KITTI_DIR> and then run:

python train.py --kitti_path <KITTI_DIR>

Further options are available to customize the training, please check them out with python train.py --help.

How to integrate NeFSAC into another RANSAC framework

Implementations and frameworks for RANSAC can vary widely, so we don't provide a ready module to use here. However, the implementation of a NeFSAC sampler from a torschscript model is straghtforward wherever torchscript is supported. A full Python implementation can be found in source/ransac/samplers.py.

Citation

If you find this work useful in your publications, please cite:

@inproceedings{cavalli2022nefsac,
  title={NeFSAC: neurally filtered minimal samples},
  author={Cavalli, Luca and Pollefeys, Marc and Barath, Daniel},
  booktitle={European Conference on Computer Vision},
  pages={351--366},
  year={2022},
  organization={Springer}
}

nefsac's Issues

a question about the training time on kitti dataset

Hello Cavalli,

Thanks for this good work.

I have a question about the training time on kitti dataset.

when I try the training pipeline with default configurations, it costs more than 5 seconds per iteration after caching the matching result. The version of CUDA I used is 11.6 with NVIDIA A100 GPU device.

The default configuration of maxlen in method main of train.py is 2, if I change the maxlen to None, it could use all of the image pairs in sequences 0, 1, 2, 3, 4 for one epoch, right? so the total number of pairs is 98784, and the training time is more than 137 hours per epoch. if we want to train more than one epoch, It takes a long time.

I wondering what are the configurations of maxlen and training epoch in your experiments.

By the way, I got two errors when running the training pipeline,

1. the first error is CUDA error: device-side assert triggered.
   when debuging, I find that there are some nans in poses of compute_rt_error. The error is gone after I add code poses = [pose for pose in poses if np.isnan(pose[0]).sum() == 0 and np.isnan(pose[1]).sum() == 0] before line 154,
   

   NeFSAC/source/data.py

   Lines 144 to 159 in f96595b

   	Es, _ = cv2.findEssentialMat(k1,
   	k2)
   	if Es is None or len(Es) == 0:
   	poses = []
   	else:
   	Es = np.reshape(Es, (Es.shape[0]//3, 3, 3))
   	def E_to_pose(e):
   	_, R, t, _ = cv2.recoverPose(e, k1, k2)
   	return R, t[:, 0]
   	poses = [E_to_pose(E) for E in Es]
   	if len(poses) > 0:
   	R = np.array([pose[0] for pose in poses])
   	T = np.array([pose[1] for pose in poses])
   	errs = rt_error(R_gt, T_gt, R, T)
   	err_gt = torch.maximum(*errs).min().item()
   	err_an = min(self.analytical_model(pose) for pose in poses)

2. the second error is name 'cache_path' is not defined
   I guess it shoud be data_cache_path in line 42 of data.py
   

   NeFSAC/source/data.py

   Lines 39 to 48 in f96595b

   	class KittiSequenceDataset(torch.utils.data.IterableDataset):
   	def __init__(self, sequence_num, bsize=256, interval=1, sample_size=5, maxlen=None, analytical='RT', kitti_base_path=None, data_cache_path=None):
   	if data_cache_path is not None:
   	self.cache = shelve.open((Path(cache_path) / f"/sq{sequence_num}_int{interval}").as_posix())
   	self.caching = True
   	print(f"Opening cache with {len(self.cache)} pre-existing entries")
   	else:
   	print("No location specified for the data cache persistency. The data for this run will not be persisted.")
   	self.cache = dict()
   	self.caching = False

plans for future updates ,about test code

Hello Cavalli,
I hope this message finds you well. Firstly, I wanted to express my appreciation for your work on NeFSAC available on GitHub. It has been a valuable resource.

I am currently working on understanding the testing phase of the code, and I am particularly interested in the implementation details related to the metrics you mentioned in your paper, specifically:

The median rotation and translation errors (in degrees)
The average run-time (in milliseconds)
The number of models tested inside USAC
I was wondering if you could share some insights or code snippets related to the computation of these metrics during the testing phase. Additionally, I noticed that the last update to the repository was 2022 . Do you have any plans for future updates, or is there any ongoing support for queries related to the code?

I understand your time is valuable, and I appreciate any guidance or assistance you can provide.

Thank you for your consideration.

A bug in sampler

Congrats on this great work!
I may found a small bug here.
See lines 19 and 26 below, I think there should be k instead of k-1. Since randint returns numbers in [low, high), currently the last correspondence will never be selected.

NeFSAC/source/ransac/samplers.py

Lines 15 to 34 in 2e74363

	def sample(self, data, *args, device, **kwargs):
	k = data.shape[0]
	# The first sample is easy to take.
	samples = torch.randint(low=0,
	high=k - 1,
	size=(self.n_iters, 1),
	device=device)
	# Now extract all the rest without repetitions.
	while samples.shape[1] < self.n_samples_per_iter:
	k = k - 1
	new_s = torch.randint(low=0,
	high=k - 1,
	size=(self.n_iters, 1),
	device=device)
	new_s = new_s + (new_s >= samples - torch.arange(
	samples.shape[1], device=device).unsqueeze(0)).long().sum(
	dim=1, keepdim=True)
	samples = torch.sort(torch.cat([samples, new_s], dim=1), dim=1)[0]
	return samples