Shortcuts

Source code for mmflow.datasets.flyingthings3d

# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
from glob import glob
from typing import Sequence, Union

from .base_dataset import BaseDataset
from .builder import DATASETS


[docs]@DATASETS.register_module() class FlyingThings3D(BaseDataset): """FlyingThings3D subset dataset. Args: direction (str): Direction of flow, has 4 options 'forward', 'backward', 'bidirection' and ['forward', 'backward']. Default: ['forward', 'backward']. scene (list, str): Scene in Flyingthings3D dataset, default: 'left'. This default value is for RAFT, as FlyingThings3D is so large and not often used, and only RAFT use the 'left' data in it. pass_style (str): Pass style for FlyingThing3D dataset, and it has 2 options ['clean', 'final']. Default: 'clean'. """ def __init__(self, *args, direction: Union[str, Sequence[str]] = ['forward', 'backward'], scene: Union[str, Sequence[str]] = 'left', pass_style: str = 'clean', **kwargs) -> None: assert direction in ['forward', 'backward', 'bidirection' ] or direction == ['forward', 'backward'] self.direction = direction scene = scene if isinstance(scene, (list, tuple)) else [scene] assert set(scene).issubset(['left', 'right']) self.scene = scene assert pass_style in ['clean', 'final'] self.pass_style = pass_style super().__init__(*args, **kwargs)
[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 = [] img2_filenames = [] flow_fw_filenames = [] flow_bw_filenames = [] flow_filenames = [] for idir, fw_dir, bw_dir in zip(self.img1_dir, self.flow_fw_dir, self.flow_bw_dir): img_filenames_ = self.get_data_filename(idir, self.img1_suffix) img_filenames_.sort() flow_fw_filenames_ = self.get_data_filename( fw_dir, self.flow_suffix) flow_fw_filenames.sort() flow_bw_filenames_ = self.get_data_filename( bw_dir, self.flow_suffix) flow_bw_filenames.sort() if self.direction == 'forward': img1_filenames += img_filenames_[:-1] img2_filenames += img_filenames_[1:] flow_filenames += flow_fw_filenames_[:-1] elif self.direction == 'backward': img1_filenames += img_filenames_[1:] img2_filenames += img_filenames_[:-1] flow_filenames += flow_bw_filenames_[1:] elif self.direction == 'bidirection': img1_filenames += img_filenames_[:-1] img2_filenames += img_filenames_[1:] flow_fw_filenames += flow_fw_filenames_[:-1] flow_bw_filenames += flow_bw_filenames_[1:] else: img1_filenames += img_filenames_[:-1] img2_filenames += img_filenames_[1:] flow_filenames += flow_fw_filenames_[:-1] img1_filenames += img_filenames_[1:] img2_filenames += img_filenames_[:-1] flow_filenames += flow_bw_filenames_[1:] self.load_img_info(self.data_infos, img1_filenames, img2_filenames) if self.direction == 'bidirection': self.load_ann_info(self.data_infos, flow_fw_filenames, 'filename_flow_fw') self.load_ann_info(self.data_infos, flow_bw_filenames, 'filename_flow_bw') else: self.load_ann_info(self.data_infos, flow_filenames, 'filename_flow')
def _get_data_dir(self) -> None: """Get the paths for images and optical flow.""" self.flow_fw_dir = 'into_future' self.flow_bw_dir = 'into_past' self.flow_suffix = '.pfm' self.occ_suffix = '.png' self.img1_suffix = '.png' self.img2_suffix = '.png' self.subset_dir = 'TEST' if self.test_mode else 'TRAIN' pass_dir = 'frames_' + self.pass_style + 'pass' imgs_dirs = [] flow_fw_dirs = [] flow_bw_dirs = [] for scene in self.scene: imgs_dirs_ = glob( osp.join(self.data_root, pass_dir, self.subset_dir + '/*/*')) imgs_dirs_ = [osp.join(f, scene) for f in imgs_dirs_] imgs_dirs += imgs_dirs_ flow_fw_dirs_ = glob( osp.join(self.data_root, 'optical_flow', self.subset_dir + '/*/*/' + self.flow_fw_dir)) flow_fw_dirs_ = [osp.join(f, scene) for f in flow_fw_dirs_] flow_fw_dirs += flow_fw_dirs_ flow_bw_dirs_ = glob( osp.join(self.data_root, 'optical_flow', self.subset_dir + '/*/*/' + self.flow_bw_dir)) flow_bw_dirs_ = [osp.join(f, scene) for f in flow_bw_dirs_] flow_bw_dirs += flow_bw_dirs_ self.img1_dir = imgs_dirs self.img2_dir = imgs_dirs self.flow_fw_dir = flow_fw_dirs self.flow_bw_dir = flow_bw_dirs
Read the Docs v: latest
Versions
latest
stable
Downloads
pdf
html
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.