PyTorch utilities to handle the KITTI Vision Benchmark Suite
MIT License
This project aims to provide a simple yet effective way to scaffold and load the KITTI Vision Banchmark Dataset providing
Datasets: Pytorch datasets to load each dataset
Scaffolding: to download the datasets
Metrics: common metrics used for each dataset
Transformations: utilities to manipulate samples
To install torch-kitti
$ pip install torch-kitti
To manually download the datasets the torch-kitti
command line utility comes in handy:
$ torch_kitti download --help
usage: Torch Kitti download [-h]
{sync_rectified,depth_completion,depth_prediction}
path
positional arguments:
{sync_rectified,depth_completion,depth_prediction}
name of the dataset to download
path where scaffold the dataset
optional arguments:
-h, --help show this help message and exit
Actually available datasets are:
All datasets return dictionaries, utilities to manipulate them can be found in the torch_kitti.transforms
module. Often each dataset provides options to include optional fields, for instance KittiDepthCompletionDataset
usually provides simply the img
, its sparse depth groundtruth gt
and the sparse lidar hints lidar
but using load_stereo=True
stereo images will be included for each example.
from torchvision.transforms import Compose, RandomCrop, ToTensor
from torch_kitti.depth_completion import KittiDepthCompletionDataset
from torch_kitti.transforms import ApplyToFeatures
transform = ApplyToFeatures(
Compose(
[
ToTensor(),
RandomCrop([256, 512]),
]
),
features=["img", "gt", "lidar"],
)
ds = KittiDepthCompletionDataset(
"kitti_raw_sync_rect_root",
"kitti_depth_completion_root",
load_stereo=False,
load_sequence=3,
transform=transform,
download=True, # download if not found
)
Each dataset exposes the elems
attribute, containing the paths used by each example loaded, and it can be modified to customize the loaded data, it is composed by DataGroup
s each containing many DataElem
s. These latter can load many different file types and each loader automatically will load them if added to each DataGroup, the supported data types are image, depth, pcd, calib, imu, rt and intrinsics.
For example to load only a specific drive sequence for depth completion you could:
class SeqKittiDepthCompletionDataset(KittiDepthCompletionDataset):
def __init__(self, drive_name: str, *args, **kwargs):
super().__init__(*args, subset="all", **kwargs)
self.elems = sorted(
filter(lambda group: group.drive == drive_name and group.cam == 2, self.elems),
key=lambda group: group.idx
)
Download from kitti and cd
in the folder then prepare a virtual environment (1), install dev
and doc
dependencies (2) and pre-commit
(3).
$ git clone https://github.com/andreaconti/torch_kitti.git
$ cd torch_kitti
$ python3 -m virtualenv .venv && source .venv/bin/activate # (1)
$ pip install .[dev, doc] # (2)
$ pre-commit install # (3)
$ python3 setup.py develop
$ pytest
Feel free to open an issue on GitHub, fork the repository and submit a pull request to solve bugs, improve docs, add datasets and features. All new feature must be tested.
This library is an utility that downloads and prepares the dataset. The KITTI Vision Benchmark Suite is not hosted by this project nor it's claimed that you have license to use the dataset, it is your responsibility to determine whether you have permission to use this dataset under its license. You can find more details here.