cityscapes.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os
  15. import glob
  16. from paddlex.paddleseg.datasets import Dataset
  17. from paddlex.paddleseg.cvlibs import manager
  18. from paddlex.paddleseg.transforms import Compose
  19. @manager.DATASETS.add_component
  20. class Cityscapes(Dataset):
  21. """
  22. Cityscapes dataset `https://www.cityscapes-dataset.com/`.
  23. The folder structure is as follow:
  24. cityscapes
  25. |
  26. |--leftImg8bit
  27. | |--train
  28. | |--val
  29. | |--test
  30. |
  31. |--gtFine
  32. | |--train
  33. | |--val
  34. | |--test
  35. Make sure there are **labelTrainIds.png in gtFine directory. If not, please run the conver_cityscapes.py in tools.
  36. Args:
  37. transforms (list): Transforms for image.
  38. dataset_root (str): Cityscapes dataset directory.
  39. mode (str, optional): Which part of dataset to use. it is one of ('train', 'val', 'test'). Default: 'train'.
  40. edge (bool, optional): Whether to compute edge while training. Default: False
  41. """
  42. NUM_CLASSES = 19
  43. def __init__(self, transforms, dataset_root, mode='train', edge=False):
  44. self.dataset_root = dataset_root
  45. self.transforms = Compose(transforms)
  46. self.file_list = list()
  47. mode = mode.lower()
  48. self.mode = mode
  49. self.num_classes = self.NUM_CLASSES
  50. self.ignore_index = 255
  51. self.edge = edge
  52. if mode not in ['train', 'val', 'test']:
  53. raise ValueError(
  54. "mode should be 'train', 'val' or 'test', but got {}.".format(
  55. mode))
  56. if self.transforms is None:
  57. raise ValueError("`transforms` is necessary, but it is None.")
  58. img_dir = os.path.join(self.dataset_root, 'leftImg8bit')
  59. label_dir = os.path.join(self.dataset_root, 'gtFine')
  60. if self.dataset_root is None or not os.path.isdir(
  61. self.dataset_root) or not os.path.isdir(
  62. img_dir) or not os.path.isdir(label_dir):
  63. raise ValueError(
  64. "The dataset is not Found or the folder structure is nonconfoumance."
  65. )
  66. label_files = sorted(
  67. glob.glob(
  68. os.path.join(label_dir, mode, '*',
  69. '*_gtFine_labelTrainIds.png')))
  70. img_files = sorted(
  71. glob.glob(os.path.join(img_dir, mode, '*', '*_leftImg8bit.png')))
  72. self.file_list = [[
  73. img_path, label_path
  74. ] for img_path, label_path in zip(img_files, label_files)]