[Feature] Support dataset initialization with file_client (#1402)

* [Feature] Support file_client in dataset initialization

* docstring

* revise osp.exist
This commit is contained in:
Miao Zheng 2022-03-28 23:53:23 +08:00 committed by GitHub
parent f3ae2342eb
commit 30864ea23d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 28 additions and 21 deletions

View File

@ -1,5 +1,4 @@
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
from .builder import DATASETS
from .custom import CustomDataset
@ -25,4 +24,4 @@ class ChaseDB1Dataset(CustomDataset):
seg_map_suffix='_1stHO.png',
reduce_zero_label=False,
**kwargs)
assert osp.exists(self.img_dir)
assert self.file_client.exists(self.img_dir)

View File

@ -68,6 +68,9 @@ class CustomDataset(Dataset):
Default: None
gt_seg_map_loader_cfg (dict, optional): build LoadAnnotations to
load gt for evaluation, load from disk by default. Default: None.
file_client_args (dict): Arguments to instantiate a FileClient.
See :class:`mmcv.fileio.FileClient` for details.
Defaults to ``dict(backend='disk')``.
"""
CLASSES = None
@ -87,7 +90,8 @@ class CustomDataset(Dataset):
reduce_zero_label=False,
classes=None,
palette=None,
gt_seg_map_loader_cfg=None):
gt_seg_map_loader_cfg=None,
file_client_args=dict(backend='disk')):
self.pipeline = Compose(pipeline)
self.img_dir = img_dir
self.img_suffix = img_suffix
@ -105,6 +109,9 @@ class CustomDataset(Dataset):
) if gt_seg_map_loader_cfg is None else LoadAnnotations(
**gt_seg_map_loader_cfg)
self.file_client_args = file_client_args
self.file_client = mmcv.FileClient.infer_client(self.file_client_args)
if test_mode:
assert self.CLASSES is not None, \
'`cls.CLASSES` or `classes` should be specified when testing'
@ -146,16 +153,21 @@ class CustomDataset(Dataset):
img_infos = []
if split is not None:
with open(split) as f:
for line in f:
img_name = line.strip()
img_info = dict(filename=img_name + img_suffix)
if ann_dir is not None:
seg_map = img_name + seg_map_suffix
img_info['ann'] = dict(seg_map=seg_map)
img_infos.append(img_info)
lines = mmcv.list_from_file(
split, file_client_args=self.file_client_args)
for line in lines:
img_name = line.strip()
img_info = dict(filename=img_name + img_suffix)
if ann_dir is not None:
seg_map = img_name + seg_map_suffix
img_info['ann'] = dict(seg_map=seg_map)
img_infos.append(img_info)
else:
for img in mmcv.scandir(img_dir, img_suffix, recursive=True):
for img in self.file_client.list_dir_or_file(
dir_path=img_dir,
list_dir=False,
suffix=img_suffix,
recursive=True):
img_info = dict(filename=img)
if ann_dir is not None:
seg_map = img.replace(img_suffix, seg_map_suffix)

View File

@ -1,5 +1,4 @@
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
from .builder import DATASETS
from .custom import CustomDataset
@ -25,4 +24,4 @@ class DRIVEDataset(CustomDataset):
seg_map_suffix='_manual1.png',
reduce_zero_label=False,
**kwargs)
assert osp.exists(self.img_dir)
assert self.file_client.exists(self.img_dir)

View File

@ -1,5 +1,4 @@
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
from .builder import DATASETS
from .custom import CustomDataset
@ -25,4 +24,4 @@ class HRFDataset(CustomDataset):
seg_map_suffix='.png',
reduce_zero_label=False,
**kwargs)
assert osp.exists(self.img_dir)
assert self.file_client.exists(self.img_dir)

View File

@ -1,5 +1,4 @@
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import mmcv
from mmcv.utils import print_log
@ -35,7 +34,7 @@ class iSAIDDataset(CustomDataset):
seg_map_suffix='.png',
ignore_index=255,
**kwargs)
assert osp.exists(self.img_dir)
assert self.file_client.exists(self.img_dir)
def load_annotations(self,
img_dir,

View File

@ -1,5 +1,4 @@
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
from .builder import DATASETS
from .custom import CustomDataset
@ -52,7 +51,7 @@ class PascalContextDataset(CustomDataset):
split=split,
reduce_zero_label=False,
**kwargs)
assert osp.exists(self.img_dir) and self.split is not None
assert self.file_client.exists(self.img_dir) and self.split is not None
@DATASETS.register_module()
@ -101,4 +100,4 @@ class PascalContextDataset59(CustomDataset):
split=split,
reduce_zero_label=True,
**kwargs)
assert osp.exists(self.img_dir) and self.split is not None
assert self.file_client.exists(self.img_dir) and self.split is not None