update pubtabnet dataset

ndcuong91 · ndcuong91 · commit ac56d5ccd10a · 2020-10-30T15:00:36.000+07:00
diff --git a/configs/_base_/datasets/publaynet.py b/configs/_base_/datasets/publaynet.py
@@ -0,0 +1,54 @@
+# dataset settings
+dataset_type = 'publaynetDataset'
+data_root = '/data4T/ntanh'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+crop_size = (512, 512)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations'),
+    dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
+    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='PhotoMetricDistortion'),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(2048, 512),
+        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=4,
+    workers_per_gpu=4,
+    train=dict(
+        type=dataset_type,
+        data_root=data_root,
+        img_dir='publaynet/train',
+        ann_dir='publaynet_gen_gt_oct2.1/train/label',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        data_root=data_root,
+        img_dir='publaynet/val',
+        ann_dir='publaynet_gen_gt_oct2.1/val/label',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        data_root=data_root,
+        img_dir='publaynet/val',
+        ann_dir='publaynet_gen_gt_oct2.1/val/label',
+        pipeline=test_pipeline))
diff --git a/configs/_base_/datasets/publaynet_split1.py b/configs/_base_/datasets/publaynet_split1.py
@@ -0,0 +1,54 @@
+# dataset settings
+dataset_type = 'publaynet_split1Dataset'
+data_root = '/data20.04/data/doc_structure/publaynet'
+img_norm_cfg = dict(
+    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
+crop_size = (512, 512)
+train_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(type='LoadAnnotations'),
+    dict(type='Resize', img_scale=(2048, 512), ratio_range=(0.5, 2.0)),
+    dict(type='RandomCrop', crop_size=crop_size, cat_max_ratio=0.75),
+    dict(type='RandomFlip', flip_ratio=0.5),
+    dict(type='PhotoMetricDistortion'),
+    dict(type='Normalize', **img_norm_cfg),
+    dict(type='Pad', size=crop_size, pad_val=0, seg_pad_val=255),
+    dict(type='DefaultFormatBundle'),
+    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
+]
+test_pipeline = [
+    dict(type='LoadImageFromFile'),
+    dict(
+        type='MultiScaleFlipAug',
+        img_scale=(2048, 512),
+        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
+        flip=False,
+        transforms=[
+            dict(type='Resize', keep_ratio=True),
+            dict(type='RandomFlip'),
+            dict(type='Normalize', **img_norm_cfg),
+            dict(type='ImageToTensor', keys=['img']),
+            dict(type='Collect', keys=['img']),
+        ])
+]
+data = dict(
+    samples_per_gpu=4,
+    workers_per_gpu=8,
+    train=dict(
+        type=dataset_type,
+        data_root=data_root,
+        img_dir='img_dir/train',
+        ann_dir='ann_dir/train',
+        pipeline=train_pipeline),
+    val=dict(
+        type=dataset_type,
+        data_root=data_root,
+        img_dir='img_dir/val',
+        ann_dir='ann_dir/val',
+        pipeline=test_pipeline),
+    test=dict(
+        type=dataset_type,
+        data_root=data_root,
+        img_dir='img_dir/val',
+        ann_dir='ann_dir/val',
+        pipeline=test_pipeline))
diff --git a/configs/_base_/datasets/table_structure1.py b/configs/_base_/datasets/table_structure1.py
@@ -1,6 +1,6 @@
 # dataset settings
 dataset_type = 'table_structure1Dataset'
-data_root = 'data/table_structure1'
+data_root = '/home/cuongnd/PycharmProjects/open-mmlab/mmsegmentation/data/table_structure1'
 img_norm_cfg = dict(
     mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
 crop_size = (512, 512)
diff --git a/configs/_base_/models/fast_scnn.py b/configs/_base_/models/fast_scnn.py
@@ -1,5 +1,5 @@
 # model settings
-norm_cfg = dict(type='SyncBN', requires_grad=True, momentum=0.01)
+norm_cfg = dict(type='BN', requires_grad=True, momentum=0.01)
 model = dict(
     type='EncoderDecoder',
     backbone=dict(
diff --git a/configs/_base_/schedules/schedule_20k.py b/configs/_base_/schedules/schedule_20k.py
@@ -5,5 +5,6 @@
 lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
 # runtime settings
 total_iters = 20000
-checkpoint_config = dict(by_epoch=False, interval=2000)
-evaluation = dict(interval=2000, metric='mIoU')
+interval=1000
+checkpoint_config = dict(by_epoch=False, interval=interval)
+evaluation = dict(interval=interval, metric='mIoU')
diff --git a/configs/_base_/schedules/schedule_320k.py b/configs/_base_/schedules/schedule_320k.py
@@ -0,0 +1,9 @@
+# optimizer
+optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0005)
+optimizer_config = dict()
+# learning policy
+lr_config = dict(policy='poly', power=0.9, min_lr=1e-5, by_epoch=False)
+# runtime settings
+total_iters = 320000
+checkpoint_config = dict(by_epoch=False, interval=32000)
+evaluation = dict(interval=320000, metric='mIoU')
diff --git a/configs/_base_/schedules/schedule_40k.py b/configs/_base_/schedules/schedule_40k.py
@@ -5,5 +5,5 @@
 lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
 # runtime settings
 total_iters = 40000
-checkpoint_config = dict(by_epoch=False, interval=4000)
-evaluation = dict(interval=4000, metric='mIoU')
+checkpoint_config = dict(by_epoch=False, interval=interval)
+evaluation = dict(interval=interval, metric='mIoU')
diff --git a/configs/_base_/schedules/schedule_80k.py b/configs/_base_/schedules/schedule_80k.py
@@ -6,4 +6,4 @@
 # runtime settings
 total_iters = 80000
 checkpoint_config = dict(by_epoch=False, interval=8000)
-evaluation = dict(interval=8000, metric='mIoU')
+evaluation = dict(interval=80000, metric='mIoU')
diff --git a/configs/fastscnn/fast_scnn_4x8_320k_lr0.12_publaynet.py b/configs/fastscnn/fast_scnn_4x8_320k_lr0.12_publaynet.py
@@ -0,0 +1,10 @@
+_base_ = [
+    '../_base_/models/fast_scnn.py', '../_base_/datasets/publaynet.py',
+    '../_base_/default_runtime.py', '../_base_/schedules/schedule_320k.py'
+]
+
+# Re-config the data sampler.
+data = dict(samples_per_gpu=16, workers_per_gpu=4)
+
+# Re-config the optimizer.
+optimizer = dict(type='SGD', lr=0.12, momentum=0.9, weight_decay=4e-5)
diff --git a/configs/fastscnn/fast_scnn_4x8_80k_lr0.12_publaynet_split1.py b/configs/fastscnn/fast_scnn_4x8_80k_lr0.12_publaynet_split1.py
@@ -0,0 +1,10 @@
+_base_ = [
+    '../_base_/models/fast_scnn.py', '../_base_/datasets/publaynet_split1.py',
+    '../_base_/default_runtime.py', '../_base_/schedules/schedule_80k.py'
+]
+
+# Re-config the data sampler.
+data = dict(samples_per_gpu=16, workers_per_gpu=4)
+
+# Re-config the optimizer.
+optimizer = dict(type='SGD', lr=0.12, momentum=0.9, weight_decay=4e-5)
diff --git a/configs/pspnet/pspnet_r50-d8_512x512_20k_publaynet.py b/configs/pspnet/pspnet_r50-d8_512x512_20k_publaynet.py
@@ -0,0 +1,8 @@
+_base_ = [
+    '../_base_/models/pspnet_r50-d8.py',
+    '../_base_/datasets/publaynet.py', '../_base_/default_runtime.py',
+    '../_base_/schedules/schedule_20k.py'
+]
+model = dict(
+    decode_head=dict(num_classes=6), auxiliary_head=dict(num_classes=6))
+test_cfg = dict(mode='whole')
diff --git a/configs/pspnet/pspnet_r50-d8_512x512_20k_table_structure1.py b/configs/pspnet/pspnet_r50-d8_512x512_20k_table_structure1.py
@@ -1,9 +1,8 @@
 _base_ = [
     '../_base_/models/pspnet_r50-d8.py',
     '../_base_/datasets/table_structure1.py', '../_base_/default_runtime.py',
-    '../_base_/schedules/schedule_40k.py'
+    '../_base_/schedules/schedule_20k.py'
 ]
 model = dict(
     decode_head=dict(num_classes=2), auxiliary_head=dict(num_classes=2))
-test_cfg = dict(mode='slide', crop_size=(480, 480), stride=(320, 320))
-optimizer = dict(type='SGD', lr=0.004, momentum=0.9, weight_decay=0.0001)
+test_cfg = dict(mode='whole')
diff --git a/configs/pspnet/pspnet_r50-d8_512x512_320k_publaynet_split1.py b/configs/pspnet/pspnet_r50-d8_512x512_320k_publaynet_split1.py
@@ -0,0 +1,8 @@
+_base_ = [
+    '../_base_/models/pspnet_r50-d8.py',
+    '../_base_/datasets/publaynet_split1.py', '../_base_/default_runtime.py',
+    '../_base_/schedules/schedule_320k.py'
+]
+model = dict(
+    decode_head=dict(num_classes=6), auxiliary_head=dict(num_classes=6))
+test_cfg = dict(mode='whole')
diff --git a/configs/pspnet/pspnet_r50-d8_512x512_40k_table_structure1.py b/configs/pspnet/pspnet_r50-d8_512x512_40k_table_structure1.py
@@ -0,0 +1,8 @@
+_base_ = [
+    '../_base_/models/pspnet_r50-d8.py',
+    '../_base_/datasets/table_structure1.py', '../_base_/default_runtime.py',
+    '../_base_/schedules/schedule_40k.py'
+]
+model = dict(
+    decode_head=dict(num_classes=1), auxiliary_head=dict(num_classes=1))
+test_cfg = dict(mode='whole')
diff --git a/configs/pspnet/pspnet_r50-d8_512x512_80k_publaynet_split1.py b/configs/pspnet/pspnet_r50-d8_512x512_80k_publaynet_split1.py
@@ -0,0 +1,8 @@
+_base_ = [
+    '../_base_/models/pspnet_r50-d8.py',
+    '../_base_/datasets/publaynet_split1.py', '../_base_/default_runtime.py',
+    '../_base_/schedules/schedule_80k.py'
+]
+model = dict(
+    decode_head=dict(num_classes=6), auxiliary_head=dict(num_classes=6))
+test_cfg = dict(mode='whole')
diff --git a/demo/image_demo.py b/demo/image_demo.py
diff --git a/mmseg/core/evaluation/class_names.py b/mmseg/core/evaluation/class_names.py
@@ -136,17 +136,20 @@ def get_classes(dataset):
 
 
 def get_palette(dataset):
-    """Get class palette (RGB) of a dataset."""
-    alias2name = {}
-    for name, aliases in dataset_aliases.items():
-        for alias in aliases:
-            alias2name[alias] = name
-
-    if mmcv.is_str(dataset):
-        if dataset in alias2name:
-            labels = eval(alias2name[dataset] + '_palette()')
+    if dataset is not None:
+        """Get class palette (RGB) of a dataset."""
+        alias2name = {}
+        for name, aliases in dataset_aliases.items():
+            for alias in aliases:
+                alias2name[alias] = name
+
+        if mmcv.is_str(dataset):
+            if dataset in alias2name:
+                labels = eval(alias2name[dataset] + '_palette()')
+            else:
+                raise ValueError(f'Unrecognized dataset: {dataset}')
         else:
-            raise ValueError(f'Unrecognized dataset: {dataset}')
+            raise TypeError(f'dataset must a str, but got {type(dataset)}')
+        return labels
     else:
-        raise TypeError(f'dataset must a str, but got {type(dataset)}')
-    return labels
+        return None
diff --git a/mmseg/datasets/__init__.py b/mmseg/datasets/__init__.py
@@ -6,9 +6,12 @@
 from .pascal_context import PascalContextDataset
 from .voc import PascalVOCDataset
 from .table_structure1 import table_structure1Dataset
+from .publaynet import publaynetDataset
+from .publaynet_split1 import publaynet_split1Dataset
 
 __all__ = [
     'CustomDataset', 'build_dataloader', 'ConcatDataset', 'RepeatDataset',
     'DATASETS', 'build_dataset', 'PIPELINES', 'CityscapesDataset',
-    'PascalVOCDataset', 'ADE20KDataset', 'PascalContextDataset', 'table_structure1Dataset'
+    'PascalVOCDataset', 'ADE20KDataset', 'PascalContextDataset', 'table_structure1Dataset',
+    'publaynetDataset','publaynet_split1Dataset'
 ]
diff --git a/mmseg/datasets/custom.py b/mmseg/datasets/custom.py
@@ -322,7 +322,6 @@ def evaluate(self, results, metric='mIoU', logger=None, **kwargs):
         Returns:
             dict[str, float]: Default metrics.
         """
-
         if not isinstance(metric, str):
             assert len(metric) == 1
             metric = metric[0]
@@ -339,7 +338,7 @@ def evaluate(self, results, metric='mIoU', logger=None, **kwargs):
             num_classes = len(self.CLASSES)
 
         all_acc, acc, iou = mean_iou(
-            results, gt_seg_maps, num_classes, ignore_index=self.ignore_index)
+            results, gt_seg_maps, num_classes, ignore_index=self.ignore_index, nan_to_num=-1)
         summary_str = ''
         summary_str += 'per class results:\n'
 
diff --git a/mmseg/datasets/publaynet.py b/mmseg/datasets/publaynet.py
@@ -0,0 +1,16 @@
+from .builder import DATASETS
+from .custom import CustomDataset
+
+
+@DATASETS.register_module()
+class publaynetDataset(CustomDataset):
+    """table_structure1
+    """
+    CLASSES = ('background','text', 'title', 'list', 'table', 'figure')
+    PALETTE = [[120, 120, 120],[50, 255, 0],[255, 0, 0],[0, 255, 255],[255, 192, 203],[100, 0, 255]]
+    def __init__(self, **kwargs):
+        super(publaynetDataset, self).__init__(
+            img_suffix='.jpg',
+            seg_map_suffix='.png',
+            reduce_zero_label=False,
+            **kwargs)
diff --git a/mmseg/datasets/publaynet_split1.py b/mmseg/datasets/publaynet_split1.py
@@ -0,0 +1,16 @@
+from .builder import DATASETS
+from .custom import CustomDataset
+
+
+@DATASETS.register_module()
+class publaynet_split1Dataset(CustomDataset):
+    """table_structure1
+    """
+    CLASSES = ('background','text', 'title', 'list', 'table', 'figure')
+    PALETTE = [[120, 120, 120],[50, 255, 0],[255, 0, 0],[0, 255, 255],[255, 192, 203],[100, 0, 255]]
+    def __init__(self, **kwargs):
+        super(publaynet_split1Dataset, self).__init__(
+            img_suffix='.jpg',
+            seg_map_suffix='.png',
+            reduce_zero_label=False,
+            **kwargs)
diff --git a/mmseg/datasets/table_structure1.py b/mmseg/datasets/table_structure1.py
@@ -6,11 +6,11 @@
 class table_structure1Dataset(CustomDataset):
     """table_structure1
     """
-    CLASSES = ('background', 'cell')
-    PALETTE = [[120, 120, 120], [6, 230, 230]]
+    CLASSES = ('background','cell')
+    PALETTE = [[120, 0, 0],[0, 120, 0]]
     def __init__(self, **kwargs):
         super(table_structure1Dataset, self).__init__(
             img_suffix='.png',
             seg_map_suffix='.png',
-            reduce_zero_label=True,
+            reduce_zero_label=False,
             **kwargs)
diff --git a/test_ckpt b/test_ckpt
@@ -1 +1,2 @@
-python tools/test.py configs/pspnet/pspnet_r50-d8_512x512_80k_ade20k.py checkpoints/pspnet_r50-d8_512x512_80k_ade20k_20200615_014128-15a8b914.pth --show
+#python tools/test.py configs/pspnet/pspnet_r50-d8_480x480_20k_table_structure1.py work_dirs/pspnet_r50-d8_480x480_20k_table_structure1/iter_2000.pth --show
+python demo/image_demo.py   work_dirs/pspnet_r50-d8_480x480_20k_table_structure1/iter_20000.pth --device cuda:0
diff --git a/tools/image_demo.py b/tools/image_demo.py
diff --git a/tools/prepare_segmentation_data.py b/tools/prepare_segmentation_data.py
diff --git a/tools/test.py b/tools/test.py
diff --git a/tools/train.py b/tools/train.py
