add custom dataset

mmseg/datasets/__init__.py

@@ -1,30 +1,30 @@
-# Copyright (c) OpenMMLab. All rights reserved.
-from .ade import ADE20KDataset
-from .builder import DATASETS, PIPELINES, build_dataloader, build_dataset
-from .chase_db1 import ChaseDB1Dataset
-from .cityscapes import CityscapesDataset
-from .coco_stuff import COCOStuffDataset
-from .custom import CustomDataset
-from .dark_zurich import DarkZurichDataset
-from .dataset_wrappers import (ConcatDataset, MultiImageMixDataset,
- RepeatDataset)
-from .drive import DRIVEDataset
-from .hrf import HRFDataset
-from .isaid import iSAIDDataset
-from .isprs import ISPRSDataset
-from .loveda import LoveDADataset
-from .night_driving import NightDrivingDataset
-from .pascal_context import PascalContextDataset, PascalContextDataset59
-from .potsdam import PotsdamDataset
-from .stare import STAREDataset
-from .voc import PascalVOCDataset
-
-__all__ = [
- 'CustomDataset', 'build_dataloader', 'ConcatDataset', 'RepeatDataset',
- 'DATASETS', 'build_dataset', 'PIPELINES', 'CityscapesDataset',
- 'PascalVOCDataset', 'ADE20KDataset', 'PascalContextDataset',
- 'PascalContextDataset59', 'ChaseDB1Dataset', 'DRIVEDataset', 'HRFDataset',
- 'STAREDataset', 'DarkZurichDataset', 'NightDrivingDataset',
- 'COCOStuffDataset', 'LoveDADataset', 'MultiImageMixDataset',
- 'iSAIDDataset', 'ISPRSDataset', 'PotsdamDataset'
-]
+# Copyright (c) OpenMMLab. All rights reserved.
+from .ade import ADE20KDataset
+from .builder import DATASETS, PIPELINES, build_dataloader, build_dataset
+from .chase_db1 import ChaseDB1Dataset
+from .cityscapes import CityscapesDataset
+from .coco_stuff import COCOStuffDataset
+from .custom import CustomDataset
+from .dark_zurich import DarkZurichDataset
+from .dataset_wrappers import (ConcatDataset, MultiImageMixDataset,
+ RepeatDataset)
+from .drive import DRIVEDataset
+from .hrf import HRFDataset
+from .isaid import iSAIDDataset
+from .isprs import ISPRSDataset
+from .loveda import LoveDADataset
+from .night_driving import NightDrivingDataset
+from .pascal_context import PascalContextDataset, PascalContextDataset59
+from .potsdam import PotsdamDataset
+from .stare import STAREDataset
+from .voc import PascalVOCDataset
+from .face import FaceOccluded
+
+__all__ = [
+ 'CustomDataset', 'build_dataloader', 'ConcatDataset', 'RepeatDataset',
+ 'DATASETS', 'build_dataset', 'PIPELINES', 'CityscapesDataset',
+ 'PascalVOCDataset', 'ADE20KDataset', 'PascalContextDataset',
+ 'PascalContextDataset59', 'ChaseDB1Dataset', 'DRIVEDataset', 'HRFDataset',
+ 'STAREDataset', 'DarkZurichDataset', 'NightDrivingDataset',
+ 'COCOStuffDataset', 'LoveDADataset', 'MultiImageMixDataset',
+ 'iSAIDDataset', 'ISPRSDataset', 'PotsdamDataset',"FaceOccluded"]

mmseg/datasets/face.py

@@ -0,0 +1,23 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import os.path as osp
+
+from .builder import DATASETS
+from .custom import CustomDataset
+
+
+@DATASETS.register_module()
+class FaceOccluded(CustomDataset):
+ """Face Occluded dataset.
+
+ Args:
+ split (str): Split txt file for Pascal VOC.
+ """
+
+ CLASSES = ('background', 'face')
+
+ PALETTE = [[0, 0, 0], [128, 0, 0]]
+
+ def __init__(self, split, **kwargs):
+ super(FaceOccluded, self).__init__(
+ img_suffix='.jpg', seg_map_suffix='.png', split=split, **kwargs)
+ assert osp.exists(self.img_dir) and self.split is not None