cocostuff.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  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 CocoStuff(Dataset):
  21. """
  22. COCO-Stuff dataset `https://github.com/nightrome/cocostuff`.
  23. The folder structure is as follow:
  24. cocostuff
  25. |
  26. |--images
  27. | |--train2017
  28. | |--val2017
  29. |
  30. |--annotations
  31. | |--train2017
  32. | |--val2017
  33. Args:
  34. transforms (list): Transforms for image.
  35. dataset_root (str): Cityscapes dataset directory.
  36. mode (str): Which part of dataset to use. it is one of ('train', 'val'). Default: 'train'.
  37. """
  38. NUM_CLASSES = 172
  39. def __init__(self, transforms, dataset_root, mode='train'):
  40. self.dataset_root = dataset_root
  41. self.transforms = Compose(transforms)
  42. self.file_list = list()
  43. mode = mode.lower()
  44. self.mode = mode
  45. self.num_classes = self.NUM_CLASSES
  46. self.ignore_index = 255
  47. if mode not in ['train', 'val']:
  48. raise ValueError(
  49. "mode should be 'train', 'val', but got {}.".format(mode))
  50. if self.transforms is None:
  51. raise ValueError("`transforms` is necessary, but it is None.")
  52. img_dir = os.path.join(self.dataset_root, 'images')
  53. label_dir = os.path.join(self.dataset_root, 'annotations')
  54. if self.dataset_root is None or not os.path.isdir(
  55. self.dataset_root) or not os.path.isdir(
  56. img_dir) or not os.path.isdir(label_dir):
  57. raise ValueError(
  58. "The dataset is not Found or the folder structure is nonconfoumance."
  59. )
  60. label_files = sorted(
  61. glob.glob(os.path.join(label_dir, mode + '2017', '*.png')))
  62. img_files = sorted(
  63. glob.glob(os.path.join(img_dir, mode + '2017', '*.jpg')))
  64. self.file_list = [[img_path, label_path] for img_path, label_path in
  65. zip(img_files, label_files)]