add create_dataset_list (#54)

* add create_dataset_list function

docs/data_prepare.md

@@ -11,7 +11,9 @@ PddleSeg已支持2种标注工具：LabelMe、精灵数据标注工具。标注
 
 最后用我们提供的数据转换脚本将上述标注工具产出的数据格式转换为模型训练时所需的数据格式。
 
-## 语义分割标注规范
+## 文件列表
+
+### 文件列表规范
 
 PaddleSeg采用通用的文件列表方式组织训练集、验证集和测试集。像素标注类别需要从0开始递增。
 
@@ -64,3 +66,81 @@ PaddleSeg采用通用的文件列表方式组织训练集、验证集和测试
 ![cityscapes_filelist](./imgs/file_list.png)
 
 完整的配置信息可以参考[`./dataset/cityscapes_demo`](../dataset/cityscapes_demo/)目录下的yaml和文件列表。
+
+### 文件列表生成
+PaddleSeg提供了生成文件列表的使用脚本，可适用于自定义数据集或cityscapes数据集，并支持通过不同的Flags来开启特定功能。
+```
+python pdseg/tools/create_dataset_list.py <your/dataset/dir> ${FLAGS}
+```
+运行后将在数据集根目录下生成训练/验证/测试集的文件列表（文件主名与`--second_folder`一致，扩展名为`.list`）。
+
+#### 命令行FLAGS列表
+
+|FLAG|用途|默认值|参数数目|
+|-|-|-|-|
+|--type|指定数据集类型，`cityscapes`或`自定义`|`自定义`|1|
+|--separator|文件列表分隔符|'&#124;'|1|
+|--folder|图片和标签集的文件夹名|'images' 'annotations'|2|
+|--second_folder|训练/验证/测试集的文件夹名|'train' 'val' 'test'|若干|
+|--format|图片和标签集的数据格式|'jpg'  'png'|2|
+|--postfix|按文件主名（无扩展名）是否包含指定后缀对图片和标签集进行筛选|''   ''（2个空字符）|2|
+
+#### 使用示例
+- **对于自定义数据集**
+
+如果用户想要生成自己数据集的文件列表，需要整理成如下的目录结构：
+```
+./dataset/   # 数据集根目录
+├── annotations      # 标注目录
+│   ├── test
+│   │   ├── ...
+│   │   └── ...
+│   ├── train
+│   │   ├── ...
+│   │   └── ...
+│   └── val
+│       ├── ...
+│       └── ...
+└── images       # 原图目录
+    ├── test
+    │   ├── ...
+    │   └── ...
+    ├── train
+    │   ├── ...
+    │   └── ...
+    └── val
+        ├── ...
+        └── ...
+Note：以上目录名可任意
+```
+必须指定自定义数据集目录，可以按需要设定FLAG。
+
+**Note:** 无需指定`--type`。
+```
+# 生成文件列表，其分隔符为空格，图片和标签集的数据格式都为png
+python pdseg/tools/create_dataset_list.py <your/dataset/dir> --separator " " --format png png
+```
+```
+# 生成文件列表，其图片和标签集的文件夹名为img和gt，训练和验证集的文件夹名为training和validation，不生成测试集列表
+python pdseg/tools/create_dataset_list.py <your/dataset/dir> \
+        --folder img gt --second_folder training validation
+```
+
+
+- **对于cityscapes数据集**
+
+必须指定cityscapes数据集目录，`--type`必须为`cityscapes`。
+
+在cityscapes类型下，部分FLAG将被重新设定，无需手动指定，具体如下：
+
+|FLAG|固定值|
+|-|-|
+|--folder|'leftImg8bit' 'gtFine'|
+|--format|'png' 'png'|
+|--postfix|'_leftImg8bit' '_gtFine_labelTrainIds'|
+
+其余FLAG可以按需要设定。
+```
+# 生成cityscapes文件列表，其分隔符为逗号
+python pdseg/tools/create_dataset_list.py <your/dataset/dir> --type cityscapes --separator ","
+```

pdseg/tools/create_dataset_list.md

-# 数据列表生成
-
-```
-python pdseg/tools/create_dataset_list.py --dataset "your/dataset/dir" --file_splitor " "
-```

pdseg/tools/create_dataset_list.py

-#!/usr/bin/python
-# -*- coding: UTF-8 -*-
+# coding: utf8
+# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import glob
 import os.path
 import argparse
 
-folder_name = {
-    'image': 'leftImg8bit',
-    'label': 'gtFine',
-}
-
-postfix = {
-    'image': '_leftImg8bit',
-    'label': '_gtFine_labelTrainIds',
-}
-
-data_format = {
-    'image': 'png',
-    'label': 'png',
-}
-
 
 def parse_args():
-    parser = argparse.ArgumentParser(description='PaddleSeg generate file list on cityscapes')
+    parser = argparse.ArgumentParser(
+        description='PaddleSeg generate file list on cityscapes or your customized dataset.')
     parser.add_argument(
-        '--dataset_root',
-        dest='dataset_root',
+        'dataset_root',
         help='dataset root directory',
-        default=None,
-        type=str)
-    parser.add_argument('--file_splitor',
-                        dest='file_splitor',
-                        help='file list splitor',
-                        default=None,
-                        type=str)
+        type=str
+    )
+    parser.add_argument(
+        '--type',
+        help='dataset type: \n'
+             '- cityscapes \n'
+             '- custom(default)',
+        default="custom",
+        type=str
+    )
+    parser.add_argument(
+        '--separator',
+        dest='separator',
+        help='file list separator',
+        default="|",
+        type=str
+    )
+    parser.add_argument(
+        '--folder',
+        help='the folder names of images and labels',
+        type=str,
+        nargs=2,
+        default=['images', 'annotations']
+    )
+    parser.add_argument(
+        '--second_folder',
+        help='the second-level folder names of train set, validation set, test set',
+        type=str,
+        nargs='*',
+        default=['train', 'val', 'test']
+    )
+    parser.add_argument(
+        '--format',
+        help='data format of images and labels, e.g. jpg or png.',
+        type=str,
+        nargs=2,
+        default=['jpg', 'png']
+    )
+    parser.add_argument(
+        '--postfix',
+        help='postfix of images or labels',
+        type=str,
+        nargs=2,
+        default=['', '']
+    )
+
     return parser.parse_args()
 
 
+def cityscape_cfg(args):
+    args.postfix = ['_leftImg8bit', '_gtFine_labelTrainIds']
+
+    args.folder = ['leftImg8bit', 'gtFine']
+
+    args.format = ['png', 'png']
+
+
 def get_files(image_or_label, dataset_split, args):
     dataset_root = args.dataset_root
-    pattern = '*%s.%s' % (postfix[image_or_label], data_format[image_or_label])
-    search_files = os.path.join(
-        dataset_root, folder_name[image_or_label], dataset_split, '*', pattern)
+    postfix = args.postfix
+    format = args.format
+    folder = args.folder
+
+    pattern = '*%s.%s' % (postfix[image_or_label], format[image_or_label])
+
+    search_files = os.path.join(dataset_root, folder[image_or_label],
+                                dataset_split, pattern)
+    search_files2 = os.path.join(dataset_root, folder[image_or_label],
+                                 dataset_split, "*", pattern)  # 包含子目录
+    search_files3 = os.path.join(dataset_root, folder[image_or_label],
+                                 dataset_split, "*", "*", pattern)  # 包含三级目录
+
     filenames = glob.glob(search_files)
+    filenames2 = glob.glob(search_files2)
+    filenames3 = glob.glob(search_files3)
+
+    filenames = filenames + filenames2 + filenames3
+
     return sorted(filenames)
 
 
-def generate_list(dataset_split, args):
+def generate_list(args):
     dataset_root = args.dataset_root
-    file_splitor = args.file_splitor
-
-    image_files = get_files('image', dataset_split, args)
-    label_files = get_files('label', dataset_split, args)
+    separator = args.separator
 
-    num_images = len(image_files)
+    for dataset_split in args.second_folder:
+        print("Creating {}.list...".format(dataset_split))
+        image_files = get_files(0, dataset_split, args)
+        label_files = get_files(1, dataset_split, args)
+        if not image_files:
+            img_dir = os.path.join(dataset_root, args.folder[0], dataset_split)
+            print("No files in {}".format(img_dir))
+            continue
+        elif not label_files:
+            label_dir = os.path.join(dataset_root, args.folder[1], dataset_split)
+            print("No files in {}".format(label_dir))
+            continue
 
-    file_list = os.path.join(dataset_root, dataset_split + '.list')
-    with open(file_list, "w") as f:
-        for item in range(num_images):
-            left = image_files[item].replace(dataset_root, '')
-            if left[0] == os.path.sep:
-                left = left.lstrip(os.path.sep)
-            right = label_files[item].replace(dataset_root, '')
-            if right[0] == os.path.sep:
-                right = right.lstrip(os.path.sep)
-            line = left + file_splitor + right + '\n'
-            f.write(line)
+        num_images = len(image_files)
+        file_list = os.path.join(dataset_root, dataset_split + '.list')
+        with open(file_list, "w") as f:
+            for item in range(num_images):
+                left = image_files[item].replace(dataset_root, '')
+                if left[0] == os.path.sep:
+                    left = left.lstrip(os.path.sep)
+                right = label_files[item].replace(dataset_root, '')
+                if right[0] == os.path.sep:
+                    right = right.lstrip(os.path.sep)
+                line = left + separator + right + '\n'
+                f.write(line)
+                print(line)
 
 
 if __name__ == '__main__':
     args = parse_args()
-    for dataset_split in ['train', 'val', 'test']:
-        generate_list(dataset_split, args)
+    if args.type == 'cityscapes':
+        cityscape_cfg(args)
+    generate_list(args)