Source code for mmflow.datasets.kiiti2012
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
from .base_dataset import BaseDataset
from .builder import DATASETS
[docs]@DATASETS.register_module()
class KITTI2012(BaseDataset):
"""KITTI flow 2012 dataset."""
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
def _get_data_dir(self) -> None:
"""Get the paths for images and optical flow."""
# only provide ground truth for training
self.subset_dir = 'training'
self.data_root = osp.join(self.data_root, self.subset_dir)
# In KITTI 2012, data in `colored_0` is original data
self.img1_dir = osp.join(self.data_root, 'colored_0')
self.img2_dir = osp.join(self.data_root, 'colored_0')
self.flow_dir = osp.join(self.data_root, 'flow_occ')
self.img1_suffix = '_10.png',
self.img2_suffix = '_11.png',
self.flow_suffix = '_10.png'
[docs] def load_data_info(self) -> None:
"""Load data information, including file path of image1, image2 and
optical flow."""
self._get_data_dir()
img1_filenames = self.get_data_filename(self.img1_dir,
self.img1_suffix)
img2_filenames = self.get_data_filename(self.img2_dir,
self.img2_suffix)
flow_filenames = self.get_data_filename(self.flow_dir,
self.flow_suffix)
self.load_img_info(self.data_infos, img1_filenames, img2_filenames)
self.load_ann_info(self.data_infos, flow_filenames, 'filename_flow')