[Feature request] First stereo model: PSMNet

🚀 Feature

Kornia already has a set of exciting deep learning based models implemented.

I would like to add to Kornia the first stereo model: PSMNet. PSMNet is used in many stereo research papers as a strong comparison baseline and this will play nice together with the kornia.geometry.camera.StereoCamera module as well.

Motivation

Getting a stereo model in Kornia will promote Kornia to the go-to library for research into modern deep learning based stereo models.

Pitch

PSMNet is already open source implemented here.

This ticket will adapt the code in that repo to work in kornia.contrib. Specefically, the implementation should look like:

class PSMNet(nn.Module):
    def __init__(self, pretrained: str,  max_disparity: int, model: str = "stackhourglass"):

    def forward(left_rectified_image_batch, right_rectified_image_batch):
    

Ideally I would also like to see a general purpose function disparity_regression that goes from a general cost volume to the disparity tensor. This is already implemented in above mentioned repository, but will be useful to have as a separate module for other stereo models.

Additional context

See the repository:

https://github.com/JiaRenChang/PSMNet

---

Consider also to contribute to Kornia universe projects 😃

With this model implemented, a nice tutorial cold be written that goes from rectified input images to disparity map to point cloud using PSMNet and the StereoCamera module.

• Tutorials: our repository containing the tutorials.

[copaah]

@ducha-aiki

Okay. I’ll research some models falling into category 1, i.e. fast SOTA stereo models that preferably generalizes well.

[ducha-aiki]

I’d say 1 is priority, some cool types of 2 is good to have, as well as original models which are close to 3.