From cdfb5d1875685206e5c636c8e46df65733917174 Mon Sep 17 00:00:00 2001 From: pedroHuang123 Date: Sat, 4 Feb 2023 11:45:39 +0800 Subject: [PATCH] add AutoFlow and CrowdFlow Datasets --- .pre-commit-config.yaml | 2 +- mmflow/datasets/AutoFlow.py | 53 ++++++++++++++++++++++++++++++++++++ mmflow/datasets/CrowdFlow.py | 52 +++++++++++++++++++++++++++++++++++ mmflow/datasets/__init__.py | 4 ++- 4 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 mmflow/datasets/AutoFlow.py create mode 100644 mmflow/datasets/CrowdFlow.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ea4f4d36..e4941290 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ repos: hooks: - id: flake8 - repo: https://github.com/PyCQA/isort - rev: 5.10.1 + rev: 5.11.5 hooks: - id: isort - repo: https://github.com/pre-commit/mirrors-yapf diff --git a/mmflow/datasets/AutoFlow.py b/mmflow/datasets/AutoFlow.py new file mode 100644 index 00000000..6c330187 --- /dev/null +++ b/mmflow/datasets/AutoFlow.py @@ -0,0 +1,53 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os +import os.path as osp +from typing import Any + +from .base_dataset import BaseDataset +from .builder import DATASETS + + +@DATASETS.register_module() +class AutoFlow(BaseDataset): + """AutoFlow dataset.""" + + def __init__(self, *args: Any, **kwargs: Any): + super().__init__(*args, **kwargs) + + def load_data_info(self) -> None: + """load data information.""" + + self.subset_dir = 'test' if self.test_mode else 'train' + + self.data_root = osp.join(self.data_root, self.subset_dir) + self.img1_dir = self.data_root + self.img2_dir = self.data_root + self.flow_root = self.data_root + + self.img_suffix = '.png' + self.flow_suffix = '.flo' + + self.all_scene = os.listdir(self.img1_dir) + + img1_filenames = [] + img2_filenames = [] + flow_filenames = [] + + for s in self.all_scene: + file_dir = os.listdir(osp.join(self.img1_dir, s)) + for i in file_dir: + img_file_dir = osp.join(self.img1_dir, s, i) + flow_file_dir = osp.join(self.flow_root, s, i) + + flow_filenames_ = self.get_data_filename( + flow_file_dir, self.flow_suffix) + img_filenames_ = self.get_data_filename( + img_file_dir, self.img_suffix) + + flow_filenames.append(flow_filenames_[0]) + img1_filenames.append(img_filenames_[0]) + img2_filenames.append(img_filenames_[1]) + + # img1_filenames, img2_filenames = self._revise_dir(flow_filenames) + self.load_img_info(self.data_infos, img1_filenames, img2_filenames) + self.load_ann_info(self.data_infos, flow_filenames, 'filename_flow') diff --git a/mmflow/datasets/CrowdFlow.py b/mmflow/datasets/CrowdFlow.py new file mode 100644 index 00000000..1492cf34 --- /dev/null +++ b/mmflow/datasets/CrowdFlow.py @@ -0,0 +1,52 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os +import os.path as osp +from typing import Any + +from .base_dataset import BaseDataset +from .builder import DATASETS + + +@DATASETS.register_module() +class CrowdFlow(BaseDataset): + """CrowdFlow dataset.""" + + def __init__(self, *args: Any, **kwargs: Any): + super().__init__(*args, **kwargs) + + def load_data_info(self) -> None: + """load data information.""" + + self.subset_dir = 'test' if self.test_mode else 'train' + + self.data_root = osp.join(self.data_root, self.subset_dir) + self.img1_dir = osp.join(self.data_root, 'images') + self.img2_dir = osp.join(self.data_root, 'images') + self.flow_root = osp.join(self.data_root, 'gt_flow') + + self.img_suffix = '.png' + self.flow_suffix = '.flo' + + self.all_scene = os.listdir(self.img1_dir) + + img1_filenames = [] + img2_filenames = [] + flow_filenames = [] + + for s in self.all_scene: + img_file_dir = osp.join(self.img1_dir, s) + flow_file_dir = osp.join(self.flow_root, s) + + flow_filenames_ = self.get_data_filename(flow_file_dir, + self.flow_suffix) + img_filenames_ = self.get_data_filename(img_file_dir, + self.img_suffix) + flow_num = len(flow_filenames_) + for i in range(flow_num): + flow_filenames += [flow_filenames_[i]] + img1_filenames += [img_filenames_[i]] + img2_filenames += [img_filenames_[i + 1]] + + # img1_filenames, img2_filenames = self._revise_dir(flow_filenames) + self.load_img_info(self.data_infos, img1_filenames, img2_filenames) + self.load_ann_info(self.data_infos, flow_filenames, 'filename_flow') diff --git a/mmflow/datasets/__init__.py b/mmflow/datasets/__init__.py index 5631e21a..4a1c53ed 100644 --- a/mmflow/datasets/__init__.py +++ b/mmflow/datasets/__init__.py @@ -1,6 +1,8 @@ # Copyright (c) OpenMMLab. All rights reserved. +from .AutoFlow import AutoFlow from .builder import DATASETS, PIPELINES, build_dataloader, build_dataset from .chairssdhom import ChairsSDHom +from .CrowdFlow import CrowdFlow from .dataset_wrappers import ConcatDataset, RepeatDataset from .flyingchairs import FlyingChairs from .flyingchairsocc import FlyingChairsOcc @@ -33,5 +35,5 @@ 'read_flow_kitti', 'GaussianNoise', 'RandomTranslate', 'Compose', 'InputPad', 'FlyingThings3DSubset', 'FlyingThings3D', 'Sintel', 'KITTI2012', 'KITTI2015', 'ChairsSDHom', 'HD1K', 'FlyingChairsOcc', - 'render_color_wheel' + 'CrowdFlow', 'AutoFlow', 'render_color_wheel' ]